use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
use serde_json::json;
use crate::application::orchestration::runtime_job_payloads::{
HandoffPatternJobPayload, HandoffTurnJobPayload,
};
use crate::application::orchestration::handoff_pattern_pipeline::{
HandoffPatternExecutionRequest, HandoffPatternPipeline, HandoffPatternTurn,
};
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 HandoffPatternJobHandler {
pipeline: HandoffPatternPipeline,
thread_store: Option<Arc<dyn ThreadStore>>,
}
impl HandoffPatternJobHandler {
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: HandoffPatternPipeline::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("handoff".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<HandoffPatternJobPayload, String> {
let payload: HandoffPatternJobPayload = serde_json::from_str(raw).map_err(|err| {
format!("policy violation: invalid handoff-pattern payload json: {err}")
})?;
if payload.initial_user_prompt.trim().is_empty() {
return Err(
"policy violation: handoff-pattern payload.initial_user_prompt must be non-empty"
.to_string(),
);
}
if payload.turns.is_empty() {
return Err(
"policy violation: handoff-pattern payload.turns must include at least one turn"
.to_string(),
);
}
for turn in &payload.turns {
Self::validate_turn(turn)?;
}
Ok(payload)
}
fn validate_turn(turn: &HandoffTurnJobPayload) -> std::result::Result<(), String> {
if turn.actor_id.trim().is_empty() {
return Err(
"policy violation: handoff-pattern payload.turns[].actor_id must be non-empty"
.to_string(),
);
}
if turn.user_prompt_template.trim().is_empty() {
return Err(
"policy violation: handoff-pattern payload.turns[].user_prompt_template must be non-empty"
.to_string(),
);
}
Ok(())
}
fn build_failure(message: String) -> JobExecutionOutcome {
let diagnostics = json!({
"provider": "stasis-orchestration-handoff",
"status": "failure",
"pattern": "handoff",
"guardrail_code": "POLICY_VIOLATION",
"policy_reason": &message,
})
.to_string();
JobExecutionOutcome::FatalFailure {
message,
execution_id: None,
diagnostics: Some(diagnostics),
}
}
}
#[async_trait]
impl JobHandler for HandoffPatternJobHandler {
fn job_type(&self) -> &'static str {
"workflow.stasis.orchestration.handoff"
}
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 HandoffPatternJobPayload {
thread_id,
initial_user_prompt,
policy_profile,
model_hint,
turns,
} = 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!("{}:handoff:start", job.id),
&thread_id,
"orchestration.handoff.started",
initial_user_prompt.clone(),
now,
)
.await;
let request = HandoffPatternExecutionRequest {
initial_user_prompt,
trace_id: Some(job.trace_id.clone()),
correlation_id: Some(job.correlation_id.clone()),
policy_profile,
model_hint,
turns: turns
.into_iter()
.map(|turn| HandoffPatternTurn {
actor_id: turn.actor_id,
user_prompt_template: turn.user_prompt_template,
system_prompt: turn.system_prompt,
policy_profile: turn.policy_profile,
model_hint: turn.model_hint,
})
.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-handoff",
"status": "failure",
"pattern": "handoff",
"error": error,
})
.to_string(),
),
});
}
};
let actor_ids: Vec<String> = response
.turns
.iter()
.map(|turn| turn.actor_id.clone())
.collect();
let handoffs: Vec<_> = response
.handoffs
.iter()
.map(|handoff| {
json!({
"from_actor_id": handoff.from_actor_id,
"to_actor_id": handoff.to_actor_id,
})
})
.collect();
self.append_thread_event(
format!("{}:handoff:completed", job.id),
&thread_id,
"orchestration.handoff.completed",
response.final_text.clone(),
Utc::now(),
)
.await;
Ok(JobExecutionOutcome::Success {
sttp_output_node_id: format!("sttp:orchestration:handoff:{}", job.id),
execution_id: None,
diagnostics: Some(
json!({
"provider": "stasis-orchestration-handoff",
"status": "success",
"pattern": "handoff",
"turns_executed": response.turns.len(),
"actor_ids": actor_ids,
"handoffs": handoffs,
"thread_id": thread_id,
"final_text": response.final_text,
"termination_reason": response.termination_reason,
})
.to_string(),
),
})
}
}