Files
herocoordinator/src/rpc.rs
Lee Smet e57c76fb36 Allow updating DAG
Signed-off-by: Lee Smet <lee.smet@hotmail.com>
2025-08-22 10:59:46 +02:00

603 lines
18 KiB
Rust

use std::{
collections::HashMap,
net::{IpAddr, SocketAddr},
sync::Arc,
};
use jsonrpsee::{
RpcModule,
server::{ServerBuilder, ServerHandle},
types::error::ErrorObjectOwned,
};
use serde::Deserialize;
use serde_json::{Value, json};
use crate::{
dag::{DagError, FlowDag, build_flow_dag},
models::{Actor, Context, Flow, Job, Message, MessageFormatType, Runner, ScriptType},
storage::RedisDriver,
time::current_timestamp,
};
pub struct AppState {
pub redis: RedisDriver,
}
impl AppState {
pub fn new(redis: RedisDriver) -> Self {
Self { redis }
}
}
// -----------------------------
// Error helpers
// -----------------------------
fn invalid_params_err<E: std::fmt::Display>(e: E) -> ErrorObjectOwned {
ErrorObjectOwned::owned(-32602, "Invalid params", Some(Value::String(e.to_string())))
}
fn storage_err(e: Box<dyn std::error::Error + Send + Sync>) -> ErrorObjectOwned {
let msg = e.to_string();
if msg.contains("Key not found") {
ErrorObjectOwned::owned(-32001, "Not Found", Some(Value::String(msg)))
} else {
ErrorObjectOwned::owned(-32010, "Storage Error", Some(Value::String(msg)))
}
}
fn dag_err(e: DagError) -> ErrorObjectOwned {
match e {
DagError::Storage(inner) => storage_err(inner),
DagError::MissingDependency { .. } => ErrorObjectOwned::owned(
-32020,
"DAG Missing Dependency",
Some(Value::String(e.to_string())),
),
DagError::CycleDetected { .. } => ErrorObjectOwned::owned(
-32021,
"DAG Cycle Detected",
Some(Value::String(e.to_string())),
),
DagError::UnknownJob { .. } => ErrorObjectOwned::owned(
-32022,
"DAG Unknown Job",
Some(Value::String(e.to_string())),
),
DagError::DependenciesIncomplete { .. } => ErrorObjectOwned::owned(
-32023,
"DAG Dependencies Incomplete",
Some(Value::String(e.to_string())),
),
DagError::FlowFailed { .. } => ErrorObjectOwned::owned(
-32024,
"DAG Flow Failed",
Some(Value::String(e.to_string())),
),
DagError::JobNotStarted { .. } => ErrorObjectOwned::owned(
-32025,
"DAG Job Not Started",
Some(Value::String(e.to_string())),
),
}
}
// -----------------------------
// Create DTOs and Param wrappers
// -----------------------------
#[derive(Debug, Deserialize)]
pub struct ActorCreate {
pub id: u32,
pub pubkey: String,
pub address: Vec<IpAddr>,
}
impl ActorCreate {
pub fn into_domain(self) -> Result<Actor, String> {
let ts = current_timestamp();
let v = json!({
"id": self.id,
"pubkey": self.pubkey,
"address": self.address,
"created_at": ts,
"updated_at": ts,
});
serde_json::from_value(v).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct ContextCreate {
pub id: u32,
pub admins: Vec<u32>,
pub readers: Vec<u32>,
pub executors: Vec<u32>,
}
impl ContextCreate {
pub fn into_domain(self) -> Result<Context, String> {
let ts = current_timestamp();
let mut v = serde_json::Map::new();
v.insert("id".to_string(), Value::from(self.id));
v.insert(
"admins".to_string(),
serde_json::to_value(self.admins).unwrap(),
);
v.insert(
"readers".to_string(),
serde_json::to_value(self.readers).unwrap(),
);
v.insert(
"executors".to_string(),
serde_json::to_value(self.executors).unwrap(),
);
v.insert("created_at".to_string(), Value::from(ts));
v.insert("updated_at".to_string(), Value::from(ts));
serde_json::from_value(Value::Object(v)).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct RunnerCreate {
pub id: u32,
pub pubkey: String,
pub address: IpAddr,
pub topic: String,
pub local: bool,
}
impl RunnerCreate {
pub fn into_domain(self) -> Result<Runner, String> {
let ts = current_timestamp();
let v = json!({
"id": self.id,
"pubkey": self.pubkey,
"address": self.address,
"topic": self.topic,
"local": self.local,
"created_at": ts,
"updated_at": ts,
});
serde_json::from_value(v).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct FlowCreate {
pub id: u32,
pub caller_id: u32,
pub context_id: u32,
pub jobs: Vec<u32>,
pub env_vars: HashMap<String, String>,
#[serde(default)]
pub result: Option<HashMap<String, String>>,
}
impl FlowCreate {
pub fn into_domain(self) -> Result<Flow, String> {
let ts = current_timestamp();
let v = json!({
"id": self.id,
"caller_id": self.caller_id,
"context_id": self.context_id,
"jobs": self.jobs,
"env_vars": self.env_vars,
"result": self.result.unwrap_or_default(),
"created_at": ts,
"updated_at": ts,
"status": "Dispatched",
});
serde_json::from_value(v).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct JobCreate {
pub id: u32,
pub caller_id: u32,
pub context_id: u32,
pub script: String,
pub script_type: ScriptType,
pub timeout: u32,
pub retries: u8,
pub env_vars: HashMap<String, String>,
#[serde(default)]
pub result: Option<HashMap<String, String>>,
pub prerequisites: Vec<String>,
pub depends: Vec<u32>,
}
impl JobCreate {
pub fn into_domain(self) -> Result<Job, String> {
let ts = current_timestamp();
let v = json!({
"id": self.id,
"caller_id": self.caller_id,
"context_id": self.context_id,
"script": self.script,
"script_type": self.script_type,
"timeout": self.timeout,
"retries": self.retries,
"env_vars": self.env_vars,
"result": self.result.unwrap_or_default(),
"prerequisites": self.prerequisites,
"depends": self.depends,
"created_at": ts,
"updated_at": ts,
"status": "Dispatched",
});
serde_json::from_value(v).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct MessageCreate {
pub id: u32,
pub caller_id: u32,
pub context_id: u32,
pub message: String,
pub message_type: ScriptType,
pub message_format_type: MessageFormatType,
pub timeout: u32,
pub timeout_ack: u32,
pub timeout_result: u32,
pub job: Vec<JobCreate>,
#[serde(default)]
pub logs: Option<Vec<String>>,
}
impl MessageCreate {
pub fn into_domain(self) -> Result<Message, String> {
let ts = current_timestamp();
let jobs: Result<Vec<Value>, String> = self
.job
.into_iter()
.map(|j| {
let jd: Job = j.into_domain()?;
serde_json::to_value(jd).map_err(|e| e.to_string())
})
.collect();
let v = json!({
"id": self.id,
"caller_id": self.caller_id,
"context_id": self.context_id,
"message": self.message,
"message_type": self.message_type,
"message_format_type": self.message_format_type, // "Html" | "Text" | "Md"
"timeout": self.timeout,
"timeout_ack": self.timeout_ack,
"timeout_result": self.timeout_result,
"job": jobs?,
"logs": self.logs.unwrap_or_default(),
"created_at": ts,
"updated_at": ts,
"status": "Dispatched",
});
serde_json::from_value(v).map_err(|e| e.to_string())
}
}
#[derive(Debug, Deserialize)]
pub struct ActorCreateParams {
pub context_id: u32,
pub actor: ActorCreate,
}
#[derive(Debug, Deserialize)]
pub struct ActorLoadParams {
pub context_id: u32,
pub id: u32,
}
#[derive(Debug, Deserialize)]
pub struct ContextCreateParams {
pub context: ContextCreate,
}
#[derive(Debug, Deserialize)]
pub struct ContextLoadParams {
pub id: u32,
}
#[derive(Debug, Deserialize)]
pub struct RunnerCreateParams {
pub context_id: u32,
pub runner: RunnerCreate,
}
#[derive(Debug, Deserialize)]
pub struct RunnerLoadParams {
pub context_id: u32,
pub id: u32,
}
#[derive(Debug, Deserialize)]
pub struct FlowCreateParams {
pub context_id: u32,
pub flow: FlowCreate,
}
#[derive(Debug, Deserialize)]
pub struct FlowLoadParams {
pub context_id: u32,
pub id: u32,
}
#[derive(Debug, Deserialize)]
pub struct JobCreateParams {
pub context_id: u32,
pub job: JobCreate,
}
#[derive(Debug, Deserialize)]
pub struct JobLoadParams {
pub context_id: u32,
pub caller_id: u32,
pub id: u32,
}
#[derive(Debug, Deserialize)]
pub struct MessageCreateParams {
pub context_id: u32,
pub message: MessageCreate,
}
#[derive(Debug, Deserialize)]
pub struct MessageLoadParams {
pub context_id: u32,
pub caller_id: u32,
pub id: u32,
}
// -----------------------------
// Rpc module builder (manual registration)
// -----------------------------
pub fn build_module(state: Arc<AppState>) -> RpcModule<()> {
let mut module: RpcModule<()> = RpcModule::new(());
// Actor
{
let state = state.clone();
module
.register_async_method("actor.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: ActorCreateParams = params.parse().map_err(invalid_params_err)?;
let actor = p.actor.into_domain().map_err(invalid_params_err)?;
state
.redis
.save_actor(p.context_id, &actor)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(actor)
}
})
.expect("register actor.create");
}
{
let state = state.clone();
module
.register_async_method("actor.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: ActorLoadParams = params.parse().map_err(invalid_params_err)?;
let actor = state
.redis
.load_actor(p.context_id, p.id)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(actor)
}
})
.expect("register actor.load");
}
// Context
{
let state = state.clone();
module
.register_async_method("context.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: ContextCreateParams = params.parse().map_err(invalid_params_err)?;
let ctx = p.context.into_domain().map_err(invalid_params_err)?;
state.redis.save_context(&ctx).await.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(ctx)
}
})
.expect("register context.create");
}
{
let state = state.clone();
module
.register_async_method("context.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: ContextLoadParams = params.parse().map_err(invalid_params_err)?;
let ctx = state.redis.load_context(p.id).await.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(ctx)
}
})
.expect("register context.load");
}
// Runner
{
let state = state.clone();
module
.register_async_method("runner.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: RunnerCreateParams = params.parse().map_err(invalid_params_err)?;
let runner = p.runner.into_domain().map_err(invalid_params_err)?;
state
.redis
.save_runner(p.context_id, &runner)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(runner)
}
})
.expect("register runner.create");
}
{
let state = state.clone();
module
.register_async_method("runner.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: RunnerLoadParams = params.parse().map_err(invalid_params_err)?;
let runner = state
.redis
.load_runner(p.context_id, p.id)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(runner)
}
})
.expect("register runner.load");
}
// Flow
{
let state = state.clone();
module
.register_async_method("flow.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: FlowCreateParams = params.parse().map_err(invalid_params_err)?;
let flow = p.flow.into_domain().map_err(invalid_params_err)?;
state
.redis
.save_flow(p.context_id, &flow)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(flow)
}
})
.expect("register flow.create");
}
{
let state = state.clone();
module
.register_async_method("flow.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: FlowLoadParams = params.parse().map_err(invalid_params_err)?;
let flow = state
.redis
.load_flow(p.context_id, p.id)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(flow)
}
})
.expect("register flow.load");
}
{
let state = state.clone();
module
.register_async_method("flow.dag", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: FlowLoadParams = params.parse().map_err(invalid_params_err)?;
let dag: FlowDag = build_flow_dag(&state.redis, p.context_id, p.id)
.await
.map_err(dag_err)?;
Ok::<_, ErrorObjectOwned>(dag)
}
})
.expect("register flow.dag");
}
// Job
{
let state = state.clone();
module
.register_async_method("job.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: JobCreateParams = params.parse().map_err(invalid_params_err)?;
let job = p.job.into_domain().map_err(invalid_params_err)?;
state
.redis
.save_job(p.context_id, &job)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(job)
}
})
.expect("register job.create");
}
{
let state = state.clone();
module
.register_async_method("job.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: JobLoadParams = params.parse().map_err(invalid_params_err)?;
let job = state
.redis
.load_job(p.context_id, p.caller_id, p.id)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(job)
}
})
.expect("register job.load");
}
// Message
{
let state = state.clone();
module
.register_async_method("message.create", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: MessageCreateParams = params.parse().map_err(invalid_params_err)?;
let message = p.message.into_domain().map_err(invalid_params_err)?;
state
.redis
.save_message(p.context_id, &message)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(message)
}
})
.expect("register message.create");
}
{
let state = state;
module
.register_async_method("message.load", move |params, _caller, _ctx| {
let state = state.clone();
async move {
let p: MessageLoadParams = params.parse().map_err(invalid_params_err)?;
let msg = state
.redis
.load_message(p.context_id, p.caller_id, p.id)
.await
.map_err(storage_err)?;
Ok::<_, ErrorObjectOwned>(msg)
}
})
.expect("register message.load");
}
module
}
// -----------------------------
// Server runners (HTTP/WS on separate listeners)
// -----------------------------
pub async fn start_http<C>(
addr: SocketAddr,
module: RpcModule<C>,
) -> Result<ServerHandle, Box<dyn std::error::Error + Send + Sync>> {
let server = ServerBuilder::default().build(addr).await?;
let handle = server.start(module);
Ok(handle)
}
pub async fn start_ws<C>(
addr: SocketAddr,
module: RpcModule<C>,
) -> Result<ServerHandle, Box<dyn std::error::Error + Send + Sync>> {
// jsonrpsee server supports both HTTP and WS; using a second listener gives us a dedicated WS port.
let server = ServerBuilder::default().build(addr).await?;
let handle = server.start(module);
Ok(handle)
}