orchestrator-llm/src/plan.rs

//! Plan structure and pure helpers.
//!
//! Mirrors `fai_hub::plan::Plan` so the JSON we emit deserializes
//! cleanly on the hub side. Plan validation lives here so it can
//! be exercised on both host (tests) and WASM targets.

use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize)]
pub struct Plan {
    pub schema_version: u32,
    #[serde(default)]
    pub goal: String,
    pub steps: Vec<PlanStep>,
}

#[derive(Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum PlanStep {
    Explain {
        text: String,
    },
    SaveFlow {
        name: String,
        yaml: String,
    },
    RunFlow {
        name: String,
        #[serde(default)]
        inputs: std::collections::HashMap<String, String>,
    },
    InstallModule {
        name: String,
        #[serde(default)]
        version: Option<String>,
    },
}

#[derive(Debug, thiserror::Error)]
pub enum PlanError {
    #[error("plan json could not be parsed: {0}")]
    Parse(String),
    #[error("plan schema_version {0} is not supported (this module emits version 1)")]
    UnsupportedSchemaVersion(u32),
    #[error("plan must contain at least one step")]
    Empty,
    #[error("step {index}: {reason}")]
    InvalidStep { index: usize, reason: String },
}

pub fn parse_and_validate(json: &str) -> Result<Plan, PlanError> {
    let plan: Plan = serde_json::from_str(json).map_err(|e| PlanError::Parse(e.to_string()))?;
    validate(&plan)?;
    Ok(plan)
}

pub fn validate(plan: &Plan) -> Result<(), PlanError> {
    if plan.schema_version != 1 {
        return Err(PlanError::UnsupportedSchemaVersion(plan.schema_version));
    }
    if plan.steps.is_empty() {
        return Err(PlanError::Empty);
    }
    for (i, step) in plan.steps.iter().enumerate() {
        let bad = match step {
            PlanStep::Explain { text } if text.is_empty() => Some("empty text"),
            PlanStep::SaveFlow { name, yaml } if name.is_empty() || yaml.is_empty() => {
                Some("empty name or yaml")
            }
            PlanStep::RunFlow { name, .. } if name.is_empty() => Some("empty name"),
            PlanStep::InstallModule { name, .. } if name.is_empty() => Some("empty name"),
            _ => None,
        };
        if let Some(reason) = bad {
            return Err(PlanError::InvalidStep {
                index: i,
                reason: reason.to_string(),
            });
        }
    }
    Ok(())
}

/// Serialize a plan to JSON. Used by both the LLM path (after
/// validating the LLM-emitted plan) and the stub fallback.
pub fn to_json(plan: &Plan) -> String {
    serde_json::to_string(plan).unwrap_or_else(|_| {
        r#"{"schema_version":1,"goal":"","steps":[{"kind":"explain","text":"plan serialization failed"}]}"#.to_string()
    })
}

/// Deterministic stub plan for the case where no LLM endpoint
/// was provided. Always emits a debug.echo flow.
pub fn build_stub_plan(goal: &str) -> String {
    let plan = Plan {
        schema_version: 1,
        goal: goal.to_string(),
        steps: vec![
            PlanStep::Explain {
                text: format!("Goal received: {goal}"),
            },
            PlanStep::Explain {
                text: "orchestrator-llm called without llm_endpoint — \
                       emitting deterministic stub plan."
                    .to_string(),
            },
            PlanStep::SaveFlow {
                name: "orchestrator-stub-greeting".to_string(),
                yaml: stub_flow_yaml(goal),
            },
        ],
    };
    to_json(&plan)
}

fn stub_flow_yaml(goal: &str) -> String {
    let safe_goal = goal.replace('\\', "\\\\").replace('"', "\\\"");
    format!(
        "schema_version: 1\n\
         name: orchestrator-stub-greeting\n\
         inputs:\n  who: text\n\
         steps:\n  - id: greet\n    use: debug.echo@^0\n    \
         with:\n      message: \"Goal: {safe_goal}; greeting from $inputs.who\"\n\
         outputs:\n  greeting: $greet.echoed\n",
    )
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
    use super::*;

    #[test]
    fn validates_minimal_plan() {
        let plan: Plan = serde_json::from_str(
            r#"{"schema_version":1,"goal":"x","steps":[{"kind":"explain","text":"hi"}]}"#,
        )
        .unwrap();
        validate(&plan).unwrap();
    }

    #[test]
    fn rejects_empty_steps() {
        let plan: Plan =
            serde_json::from_str(r#"{"schema_version":1,"goal":"x","steps":[]}"#).unwrap();
        assert!(matches!(validate(&plan), Err(PlanError::Empty)));
    }

    #[test]
    fn rejects_unsupported_schema_version() {
        let plan: Plan = serde_json::from_str(
            r#"{"schema_version":99,"goal":"x","steps":[{"kind":"explain","text":"hi"}]}"#,
        )
        .unwrap();
        assert!(matches!(
            validate(&plan),
            Err(PlanError::UnsupportedSchemaVersion(99))
        ));
    }

    #[test]
    fn rejects_save_flow_with_empty_name() {
        let plan: Plan = serde_json::from_str(
            r#"{"schema_version":1,"goal":"x","steps":[{"kind":"save_flow","name":"","yaml":"a"}]}"#,
        )
        .unwrap();
        assert!(matches!(
            validate(&plan),
            Err(PlanError::InvalidStep { .. })
        ));
    }

    #[test]
    fn round_trips_install_module_with_version() {
        let json = r#"{"schema_version":1,"goal":"x","steps":[{"kind":"install_module","name":"text.extract","version":"^1"}]}"#;
        let plan = parse_and_validate(json).unwrap();
        match &plan.steps[0] {
            PlanStep::InstallModule { name, version } => {
                assert_eq!(name, "text.extract");
                assert_eq!(version.as_deref(), Some("^1"));
            }
            _ => panic!("wrong variant"),
        }
    }

    #[test]
    fn stub_plan_validates() {
        let json = build_stub_plan("test goal");
        let plan: Plan = serde_json::from_str(&json).unwrap();
        validate(&plan).unwrap();
    }
}