use crate::core::engine::Engine;
use crate::core::state::{DispatchOutcome, TaskSpec};
use crate::store::run::{RunContext, StepEntry};
use crate::types::{now_unix, CapToken};
use crate::worker::adapter::SpawnerAdapter;
use async_trait::async_trait;
pub mod compiler;
pub mod loader;
pub mod store;
use mlua_flow_ir::{AsyncDispatcher, EvalError};
use serde_json::Value;
use std::sync::Arc;
pub use mlua_swarm_schema::OperatorKind as SchemaOperatorKind;
pub use mlua_swarm_schema::{
current_schema_version, default_global_agent_kind, AgentDef, AgentKind, AgentMeta,
AgentProfile, Blueprint, BlueprintMetadata, BlueprintOrigin, CompilerHints, CompilerStrategy,
OperatorDef, SpawnerHints, CURRENT_SCHEMA_VERSION,
};
pub struct EngineDispatcher {
engine: Engine,
op_token: CapToken,
spawner: Arc<dyn SpawnerAdapter>,
run_ctx: Option<RunContext>,
}
impl EngineDispatcher {
pub fn with_spawner(
engine: Engine,
op_token: CapToken,
spawner: Arc<dyn SpawnerAdapter>,
) -> Self {
Self {
engine,
op_token,
spawner,
run_ctx: None,
}
}
pub fn with_run(mut self, run_ctx: RunContext) -> Self {
self.run_ctx = Some(run_ctx);
self
}
}
#[async_trait]
impl AsyncDispatcher for EngineDispatcher {
async fn dispatch(&self, ref_: &str, input: Value) -> Result<Value, EvalError> {
let directive = match &input {
Value::String(s) => s.clone(),
other => other.to_string(),
};
let tid = self
.engine
.start_task(
&self.op_token,
TaskSpec {
agent: ref_.to_string(),
initial_directive: directive,
},
)
.await
.map_err(|e| EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("start_task: {e}"),
})?;
let run_id_for_ctx = self.run_ctx.as_ref().map(|rc| rc.run_id.clone());
let outcome = self
.engine
.dispatch_attempt_with(&self.op_token, &tid, &self.spawner, run_id_for_ctx.as_ref())
.await;
if let Some(rc) = &self.run_ctx {
let status = match &outcome {
Ok(DispatchOutcome::Pass(_)) => "passed",
Ok(DispatchOutcome::Blocked(_)) => "blocked",
Ok(DispatchOutcome::Suspended(_)) => "suspended",
Ok(DispatchOutcome::Cancelled) => "cancelled",
Ok(DispatchOutcome::Timeout) => "timeout",
Err(_) => "failed",
};
let entry = StepEntry {
step_id: tid.clone(),
step_ref: Some(ref_.to_string()),
status: Some(status.to_string()),
at: now_unix(),
};
if let Err(e) = rc.run_store.append_step_entry(&rc.run_id, entry).await {
tracing::warn!(
run_id = %rc.run_id,
step_id = %tid,
error = %e,
"EngineDispatcher::dispatch: append_step_entry failed"
);
}
}
match outcome {
Ok(DispatchOutcome::Pass(v)) => Ok(v),
Ok(DispatchOutcome::Blocked(v)) => Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("blocked: {v}"),
}),
Ok(other) => Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("non-terminal outcome: {:?}", other),
}),
Err(e) => Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("dispatch_attempt: {e}"),
}),
}
}
}