From 84d75b4edaee8eee32cc9b66651e4c841f98150e Mon Sep 17 00:00:00 2001
From: Sawyer Bergeron <sawyerbergeron@gmail.com>
Date: Wed, 26 Jun 2024 15:36:27 -0400
Subject: [PATCH] Add schema generation to tasks

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>
---
 src/Cargo.lock               | 1 +
 src/tascii/Cargo.toml        | 1 +
 src/tascii/src/task_shim.rs  | 7 +++++++
 src/tascii/src/task_trait.rs | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/src/Cargo.lock b/src/Cargo.lock
index 9ade357..ded044a 100644
--- a/src/Cargo.lock
+++ b/src/Cargo.lock
@@ -3272,6 +3272,7 @@ dependencies = [
  "pyo3",
  "rand",
  "rayon",
+ "schemars",
  "serde",
  "serde_json",
  "timer",
diff --git a/src/tascii/Cargo.toml b/src/tascii/Cargo.toml
index 1474999..35b7bce 100644
--- a/src/tascii/Cargo.toml
+++ b/src/tascii/Cargo.toml
@@ -28,6 +28,7 @@ timer = "0.2.0"
 execute = "0.2.11"
 derive_more = "0.99.17"
 itertools = "0.10.5"
+schemars = "*"
 
 tokio = { version = "1", features = [
   "full",
diff --git a/src/tascii/src/task_shim.rs b/src/tascii/src/task_shim.rs
index 100fe7f..37fd333 100644
--- a/src/tascii/src/task_shim.rs
+++ b/src/tascii/src/task_shim.rs
@@ -7,6 +7,7 @@ use std::{
 
 use dal::ID;
 
+use schemars::{schema::RootSchema, schema_for, JsonSchema};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use std::sync::Arc;
 use tracing::{debug, info};
@@ -55,6 +56,8 @@ pub(crate) trait DynRunnable: Send + std::fmt::Debug + Sync {
 
     fn unmarshal(&self, h: SimpleOneshotHandle) -> StrongUntypedOneshotHandle;
 
+    fn schema(&self) -> RootSchema;
+
     fn complete_with(
         &self,
         oneshot: &StrongUntypedOneshotHandle,
@@ -189,6 +192,10 @@ where
         });
     }
 
+    fn schema(&self) -> RootSchema {
+        schema_for!(R)
+    }
+
     fn on_error(&self) -> Box<dyn Fn(StrongUntypedOneshotHandle, TaskError) -> Result<(), ()>> {
         let summary = self.summarize(ID::nil());
         fn on_error<Res: TaskSafe>(h: StrongUntypedOneshotHandle, e: TaskError) -> Result<(), ()> {
diff --git a/src/tascii/src/task_trait.rs b/src/tascii/src/task_trait.rs
index 75e0393..d12b0ec 100644
--- a/src/tascii/src/task_trait.rs
+++ b/src/tascii/src/task_trait.rs
@@ -5,6 +5,7 @@ use std::{hash::Hash, panic::RefUnwindSafe, time::Duration, any::type_name};
 
 use dal::ID;
 
+use schemars::JsonSchema;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::{
@@ -25,6 +26,7 @@ pub trait Runnable:
     + TaskRegistered
     + 'static
     + RefUnwindSafe
+    + JsonSchema
 {
     type Output: TaskSafe;
 
@@ -78,6 +80,7 @@ pub trait AsyncRunnable:
     + TaskRegistered
     + 'static
     + RefUnwindSafe
+    + JsonSchema
 {
     type Output: TaskSafe;