cellstate_pipeline/pack/

mod.rs

//! Pack compiler -- the canonical entry point for compiling a CELLSTATE pack.
//!
//! Call [`compose_pack()`] with a [`PackInput`] (TOML manifest + Markdown files) to
//! get a [`PackOutput`] containing both the intermediate `CellstateAst` and the final
//! `CompiledConfig`.
//!
//! This module orchestrates the full pipeline:
//!
//! 1. Parse the TOML manifest (`cstate.toml`) via [`schema::parse_manifest`].
//! 2. Parse Markdown prompt files, extracting fence blocks ([`markdown`] sub-module).
//! 3. Build the intermediate representation ([`ir::PackIr`]), merging TOML and
//!    Markdown-derived definitions and validating cross-references.
//! 4. Convert the IR into a [`CellstateAst`] ([`ast::build_ast`]).
//! 5. Compile the AST into a [`CompiledConfig`] via [`PipelineCompiler::compile`].
//! 6. Inject pack-specific metadata (tools, toolsets, agents, injections, routing,
//!    file hashes) into the compiled config.
//!
//! The `ast` module provides the shared types that this module produces. The
//! `config/` module provides the YAML fence-block parsers that this module calls
//! during step 2.

mod ast;
pub mod flow;
mod ir;
mod markdown;
mod mcp;
mod schema;

// AAIF convergence modules — emit/consume open agent ecosystem formats
pub mod a2a;
pub mod a2ui_compat;
pub mod agents_md;
pub mod llms_txt;
pub mod skills;

use crate::{
    ast::CellstateAst,
    compiler::{CompiledConfig, ComposioGatewayConfig},
    PipelineCompiler,
};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::PathBuf;

pub use ast::build_ast;
pub use flow::*;
pub use ir::*;
pub use markdown::*;
pub use mcp::*;
pub use schema::*;

#[derive(Debug, Clone)]
pub struct PackInput {
    pub root: PathBuf,
    pub manifest: String,
    pub markdowns: Vec<PackMarkdownFile>,
    /// Contract JSON Schema files: maps relative path to content.
    pub contracts: HashMap<String, String>,
}

#[derive(Debug, Clone)]
pub struct PackMarkdownFile {
    pub path: PathBuf,
    pub content: String,
}

#[derive(Debug, Clone)]
pub struct PackOutput {
    pub ast: CellstateAst,
    pub compiled: CompiledConfig,
}

pub fn compose_pack(input: PackInput) -> Result<PackOutput, PackError> {
    let manifest = parse_manifest(&input.manifest)?;

    // Enforce the x-* extension namespace. Any top-level key not matching a
    // known field AND not prefixed with "x-" is a typo or unsupported feature.
    for key in manifest.extensions.keys() {
        if !key.starts_with("x-") {
            return Err(PackError::Validation(format!(
                "Unknown top-level key '{}'. Custom extensions must use the 'x-' prefix (e.g. 'x-{}')",
                key, key
            )));
        }
    }

    let md_docs = parse_markdown_files(&manifest, &input.markdowns)?;
    let ir = PackIr::new(manifest.clone(), md_docs)?;
    let ast = build_ast(&ir)?;
    let mut compiled = PipelineCompiler::compile(&ast)?;

    // Inject pack tool registry + toolsets into the runtime config.
    // Pass contract files for schema compilation.
    compiled.tools = compile_tools(&manifest, &input.contracts)?;
    compiled.toolsets = compile_toolsets(&manifest);
    compiled.pack_agents = compile_pack_agents(&manifest, &ir.markdown);
    compiled.pack_injections = compile_pack_injections(&manifest)?;
    compiled.pack_routing = compile_pack_routing(&manifest);
    compiled.composio_gateway = compile_composio_gateway(&manifest);

    // Compute file hashes for artifact determinism (lockfile support).
    compiled.file_hashes = compute_file_hashes(&input)?;

    // Compile flow definitions from Markdown fence blocks.
    compiled.flows = compile_flows(&ir.flows)?;

    // Inject pack metadata for AAIF convergence outputs.
    inject_pack_metadata(&manifest, &mut compiled);

    Ok(PackOutput { ast, compiled })
}

/// Compile flow definitions into compiled flows with content hashes.
fn compile_flows(flows: &[flow::FlowDefinition]) -> Result<Vec<flow::CompiledFlow>, PackError> {
    flows
        .iter()
        .map(|flow_def| {
            let steps: Vec<flow::CompiledFlowStep> = flow_def
                .steps
                .iter()
                .map(|step| flow::CompiledFlowStep {
                    id: step.id.clone(),
                    tool: step.tool.clone(),
                    inputs: step.inputs.clone(),
                    outputs: step.outputs.clone(),
                    content_hash: step.compute_hash(),
                })
                .collect();

            // Compute flow-level hash from name + all step hashes
            let mut hasher = Sha256::new();
            hasher.update(flow_def.name.as_bytes());
            for step in &steps {
                hasher.update(step.content_hash);
            }
            let flow_hash: [u8; 32] = hasher.finalize().into();

            Ok(flow::CompiledFlow {
                name: flow_def.name.clone(),
                flow_hash,
                steps,
            })
        })
        .collect()
}

/// Inject pack-level metadata from the manifest into CompiledConfig.
///
/// These fields power the AAIF convergence outputs: SKILL.md, AGENTS.md,
/// A2A Agent Card, llms.txt, and Google A2UI compatibility.
fn inject_pack_metadata(manifest: &PackManifest, compiled: &mut CompiledConfig) {
    if let Some(ref meta) = manifest.meta {
        compiled.pack_meta_project = meta.project.clone();
        compiled.pack_meta_version = meta.version.clone();
        compiled.pack_meta_description = meta.description.clone();
        compiled.pack_meta_instructions = meta.instructions.clone();
        compiled.pack_meta_homepage = meta.homepage.clone();
        compiled.pack_meta_license = meta.license.clone();
    }
}

/// Compile Composio MCP Gateway configuration from manifest.
fn compile_composio_gateway(manifest: &PackManifest) -> Option<ComposioGatewayConfig> {
    manifest
        .tools
        .composio_gateway
        .as_ref()
        .filter(|gw| gw.enabled)
        .map(|gw| ComposioGatewayConfig {
            enabled: gw.enabled,
            max_tools: gw.max_tools,
            allowed_toolkits: gw.allowed_toolkits.clone(),
            blocked_toolkits: gw.blocked_toolkits.clone(),
        })
}

/// Compute SHA-256 hashes for all pack source files.
/// These hashes enable runtime drift detection without recompilation.
fn compute_file_hashes(input: &PackInput) -> Result<HashMap<String, String>, PackError> {
    let mut hashes = HashMap::new();

    // Hash the manifest (cstate.toml)
    let manifest_hash = sha256_hex(&input.manifest);
    hashes.insert("cstate.toml".to_string(), manifest_hash);

    // Hash all markdown files
    for md in &input.markdowns {
        let rel_path = md
            .path
            .strip_prefix(&input.root)
            .unwrap_or(&md.path)
            .to_string_lossy()
            .to_string();
        let hash = sha256_hex(&md.content);
        hashes.insert(rel_path, hash);
    }

    // Hash all contract files
    for (path, content) in &input.contracts {
        if path.contains("..") || path.starts_with('/') || path.starts_with('\\') {
            return Err(PackError::Validation(format!(
                "invalid contract path: {path}"
            )));
        }
        let hash = sha256_hex(content);
        hashes.insert(path.clone(), hash);
    }

    Ok(hashes)
}

/// Compute SHA-256 hash of content, returning hex-encoded string.
fn sha256_hex(content: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(content.as_bytes());
    let result = hasher.finalize();
    hex::encode(result)
}

#[cfg(test)]
mod security_tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn test_contract_path_traversal_rejected() {
        let mut contracts = HashMap::new();
        contracts.insert("../etc/passwd".to_string(), "malicious".to_string());

        let input = PackInput {
            root: PathBuf::from("/tmp/test"),
            manifest: String::new(),
            markdowns: vec![],
            contracts,
        };

        let result = compute_file_hashes(&input);
        assert!(result.is_err(), "paths with '..' must be rejected");
        let err = result.unwrap_err();
        match &err {
            PackError::Validation(msg) => {
                assert!(
                    msg.contains("invalid contract path"),
                    "error should mention invalid path: {msg}"
                );
            }
            other => panic!("expected PackError::Validation, got {:?}", other),
        }

        // Also test absolute paths
        let mut contracts2 = HashMap::new();
        contracts2.insert("/etc/passwd".to_string(), "malicious".to_string());
        let input2 = PackInput {
            root: PathBuf::from("/tmp/test"),
            manifest: String::new(),
            markdowns: vec![],
            contracts: contracts2,
        };
        assert!(
            compute_file_hashes(&input2).is_err(),
            "absolute paths must be rejected"
        );

        // Test backslash-prefixed paths
        let mut contracts3 = HashMap::new();
        contracts3.insert("\\windows\\system32".to_string(), "malicious".to_string());
        let input3 = PackInput {
            root: PathBuf::from("/tmp/test"),
            manifest: String::new(),
            markdowns: vec![],
            contracts: contracts3,
        };
        assert!(
            compute_file_hashes(&input3).is_err(),
            "backslash-prefixed paths must be rejected"
        );

        // Test valid paths still work
        let mut valid_contracts = HashMap::new();
        valid_contracts.insert("schemas/contract.json".to_string(), "{}".to_string());
        let valid_input = PackInput {
            root: PathBuf::from("/tmp/test"),
            manifest: String::new(),
            markdowns: vec![],
            contracts: valid_contracts,
        };
        assert!(
            compute_file_hashes(&valid_input).is_ok(),
            "valid relative paths should be accepted"
        );
    }
}

#[cfg(test)]
mod flow_tests {
    use super::*;
    use crate::pack::flow::{FlowDefinition, FlowErrorHandler, FlowStep};
    use std::collections::HashMap;

    #[test]
    fn test_compile_flows_empty() {
        let result = compile_flows(&[]);
        assert!(result.is_ok());
        assert!(result.unwrap().is_empty());
    }

    #[test]
    fn test_compile_flows_single() {
        let flow = FlowDefinition {
            name: "test-flow".to_string(),
            description: Some("A test flow".to_string()),
            steps: vec![FlowStep {
                id: "step1".to_string(),
                tool: "bash".to_string(),
                inputs: {
                    let mut m = HashMap::new();
                    m.insert("command".to_string(), "echo hello".to_string());
                    m
                },
                outputs: vec!["stdout".to_string()],
                content_hash: [0u8; 32], // Will be recomputed
            }],
            on_error: FlowErrorHandler::Abort,
        };

        let result = compile_flows(&[flow]).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].name, "test-flow");
        assert_eq!(result[0].steps.len(), 1);
        assert_eq!(result[0].steps[0].tool, "bash");
        // Verify hash was computed (not all zeros)
        assert_ne!(result[0].flow_hash, [0u8; 32]);
        assert_ne!(result[0].steps[0].content_hash, [0u8; 32]);
    }

    #[test]
    fn test_compile_flows_multiple() {
        let flow1 = FlowDefinition {
            name: "flow-a".to_string(),
            description: None,
            steps: vec![FlowStep {
                id: "s1".to_string(),
                tool: "http".to_string(),
                inputs: HashMap::new(),
                outputs: vec![],
                content_hash: [0u8; 32],
            }],
            on_error: FlowErrorHandler::Abort,
        };

        let flow2 = FlowDefinition {
            name: "flow-b".to_string(),
            description: Some("second flow".to_string()),
            steps: vec![
                FlowStep {
                    id: "s1".to_string(),
                    tool: "bash".to_string(),
                    inputs: {
                        let mut m = HashMap::new();
                        m.insert("cmd".to_string(), "ls".to_string());
                        m
                    },
                    outputs: vec!["listing".to_string()],
                    content_hash: [0u8; 32],
                },
                FlowStep {
                    id: "s2".to_string(),
                    tool: "file_write".to_string(),
                    inputs: {
                        let mut m = HashMap::new();
                        m.insert("path".to_string(), "/tmp/out.txt".to_string());
                        m
                    },
                    outputs: vec![],
                    content_hash: [0u8; 32],
                },
            ],
            on_error: FlowErrorHandler::SkipToNext,
        };

        let result = compile_flows(&[flow1, flow2]).unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].name, "flow-a");
        assert_eq!(result[1].name, "flow-b");
        assert_eq!(result[1].steps.len(), 2);
        // Different flows should have different flow hashes
        assert_ne!(result[0].flow_hash, result[1].flow_hash);
    }

    #[test]
    fn test_compile_flows_deterministic_hashes() {
        let make_flow = || FlowDefinition {
            name: "deterministic".to_string(),
            description: None,
            steps: vec![FlowStep {
                id: "step1".to_string(),
                tool: "echo".to_string(),
                inputs: {
                    let mut m = HashMap::new();
                    m.insert("msg".to_string(), "hello".to_string());
                    m
                },
                outputs: vec!["out".to_string()],
                content_hash: [0u8; 32],
            }],
            on_error: FlowErrorHandler::Abort,
        };

        let result1 = compile_flows(&[make_flow()]).unwrap();
        let result2 = compile_flows(&[make_flow()]).unwrap();

        assert_eq!(result1[0].flow_hash, result2[0].flow_hash);
        assert_eq!(
            result1[0].steps[0].content_hash,
            result2[0].steps[0].content_hash
        );
    }
}