use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
use serde_json::json;
use crate::application::orchestration::runtime_job_payloads::{
OrchestratorPatternJobPayload, OrchestratorRouteJobPayload,
};
use crate::application::runtime::chat_options_resolver::validate_reasoning_effort;
use crate::application::orchestration::orchestrator_pattern_pipeline::{
OrchestratorPatternExecutionRequest, OrchestratorPatternPipeline, OrchestratorPatternRoute,
};
use crate::application::orchestration::prompt_pipeline::PromptExecutionPipeline;
use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
use crate::domain::errors::Result;
use crate::domain::runtime::job::Job;
use crate::domain::runtime::thread::{NewThread, NewThreadEvent};
use crate::ports::outbound::ai_chat_client::AiChatClient;
use crate::ports::outbound::runtime::thread_store::ThreadStore;
pub struct OrchestratorPatternJobHandler {
pipeline: OrchestratorPatternPipeline,
thread_store: Option<Arc<dyn ThreadStore>>,
}
impl OrchestratorPatternJobHandler {
pub fn new(chat_client: Arc<dyn AiChatClient>) -> Self {
Self::new_with_thread_store(chat_client, None)
}
pub fn new_with_thread_store(
chat_client: Arc<dyn AiChatClient>,
thread_store: Option<Arc<dyn ThreadStore>>,
) -> Self {
let prompt_pipeline = PromptExecutionPipeline::new(chat_client);
Self {
pipeline: OrchestratorPatternPipeline::new(prompt_pipeline),
thread_store,
}
}
async fn ensure_thread(&self, thread_id: &str, now: chrono::DateTime<Utc>) {
let Some(store) = &self.thread_store else {
return;
};
let exists = store.get_thread(thread_id).await.ok().flatten().is_some();
if exists {
return;
}
let _ = store
.create_thread(NewThread {
thread_id: thread_id.to_string(),
parent_thread_id: None,
branch_label: Some("orchestrator".to_string()),
created_at: now,
})
.await;
}
async fn append_thread_event(
&self,
event_id: String,
thread_id: &str,
event_kind: &str,
payload_ref: String,
occurred_at: chrono::DateTime<Utc>,
) {
let Some(store) = &self.thread_store else {
return;
};
let _ = store
.append_event(NewThreadEvent {
event_id,
thread_id: thread_id.to_string(),
event_kind: event_kind.to_string(),
payload_ref,
occurred_at,
})
.await;
}
fn parse_payload(raw: &str) -> std::result::Result<OrchestratorPatternJobPayload, String> {
let payload: OrchestratorPatternJobPayload = serde_json::from_str(raw).map_err(|err| {
format!("policy violation: invalid orchestrator-pattern payload json: {err}")
})?;
if payload.initial_user_prompt.trim().is_empty() {
return Err(
"policy violation: orchestrator-pattern payload.initial_user_prompt must be non-empty"
.to_string(),
);
}
if payload.routes.is_empty() {
return Err(
"policy violation: orchestrator-pattern payload.routes must include at least one route"
.to_string(),
);
}
for route in &payload.routes {
Self::validate_route(route)?;
}
validate_reasoning_effort(payload.reasoning_effort.as_deref())
.map_err(|err| format!("policy violation: {err}"))?;
Ok(payload)
}
fn validate_route(route: &OrchestratorRouteJobPayload) -> std::result::Result<(), String> {
if route.route_id.trim().is_empty() {
return Err(
"policy violation: orchestrator-pattern payload.routes[].route_id must be non-empty"
.to_string(),
);
}
if route.user_prompt_template.trim().is_empty() {
return Err(
"policy violation: orchestrator-pattern payload.routes[].user_prompt_template must be non-empty"
.to_string(),
);
}
validate_reasoning_effort(route.reasoning_effort.as_deref())
.map_err(|err| format!("policy violation: {err}"))?;
Ok(())
}
fn build_failure(message: String) -> JobExecutionOutcome {
let diagnostics = json!({
"provider": "stasis-orchestration-orchestrator",
"status": "failure",
"pattern": "orchestrator",
"guardrail_code": "POLICY_VIOLATION",
"policy_reason": &message,
})
.to_string();
JobExecutionOutcome::FatalFailure {
message,
execution_id: None,
diagnostics: Some(diagnostics),
}
}
}
#[async_trait]
impl JobHandler for OrchestratorPatternJobHandler {
fn job_type(&self) -> &'static str {
"workflow.stasis.orchestration.orchestrator"
}
async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
let payload = match Self::parse_payload(&job.payload_ref) {
Ok(payload) => payload,
Err(message) => return Ok(Self::build_failure(message)),
};
let OrchestratorPatternJobPayload {
thread_id,
initial_user_prompt,
policy_profile,
model_hint,
reasoning_effort,
routes,
} = payload;
let now = Utc::now();
let thread_id = thread_id.unwrap_or_else(|| job.correlation_id.clone());
self.ensure_thread(&thread_id, now).await;
self.append_thread_event(
format!("{}:orchestrator:start", job.id),
&thread_id,
"orchestration.orchestrator.started",
initial_user_prompt.clone(),
now,
)
.await;
let request = OrchestratorPatternExecutionRequest {
initial_user_prompt,
trace_id: Some(job.trace_id.clone()),
correlation_id: Some(job.correlation_id.clone()),
policy_profile,
model_hint,
reasoning_effort,
routes: routes
.into_iter()
.map(|route| OrchestratorPatternRoute {
route_id: route.route_id,
selector_keywords: route.selector_keywords,
user_prompt_template: route.user_prompt_template,
system_prompt: route.system_prompt,
policy_profile: route.policy_profile,
model_hint: route.model_hint,
reasoning_effort: route.reasoning_effort,
})
.collect(),
};
let response = match self.pipeline.execute(request).await {
Ok(response) => response,
Err(err) => {
let error = err.to_string();
return Ok(JobExecutionOutcome::FatalFailure {
message: error.clone(),
execution_id: None,
diagnostics: Some(
json!({
"provider": "stasis-orchestration-orchestrator",
"status": "failure",
"pattern": "orchestrator",
"error": error,
})
.to_string(),
),
});
}
};
self.append_thread_event(
format!("{}:orchestrator:selected", job.id),
&thread_id,
"orchestration.orchestrator.completed",
format!(
"route={} reason={}",
response.selected_route_id, response.selection_reason
),
Utc::now(),
)
.await;
Ok(JobExecutionOutcome::Success {
sttp_output_node_id: format!("sttp:orchestration:orchestrator:{}", job.id),
execution_id: None,
diagnostics: Some(
json!({
"provider": "stasis-orchestration-orchestrator",
"status": "success",
"pattern": "orchestrator",
"thread_id": thread_id,
"selected_route_id": response.selected_route_id,
"selection_reason": response.selection_reason,
"rendered_prompt": response.rendered_prompt,
"final_text": response.output_text,
"termination_reason": response.termination_reason,
})
.to_string(),
),
})
}
}