stasis/application/runtime/

orchestrator_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::{
    OrchestratorPatternJobPayload, OrchestratorRouteJobPayload,
};
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)?;
        }

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

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