stasis/application/runtime/

handoff_pattern_job_handler.rs

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(),
            ),
        })
    }
}