crank/crates/mcpaas-runtime/src/executor.rs

use std::collections::BTreeMap;

use mcpaas_adapter_rest::{RestAdapter, RestRequest};
use mcpaas_core::Target;
use serde_json::{Map, Value, json};

use crate::{AdapterResponse, PreparedRequest, RuntimeError, RuntimeOperation};

#[derive(Clone, Debug)]
pub struct RuntimeExecutor {
    rest_adapter: RestAdapter,
}

impl Default for RuntimeExecutor {
    fn default() -> Self {
        Self::new()
    }
}

impl RuntimeExecutor {
    pub fn new() -> Self {
        Self {
            rest_adapter: RestAdapter::new(),
        }
    }

    pub async fn execute(
        &self,
        operation: &RuntimeOperation,
        input: &Value,
    ) -> Result<Value, RuntimeError> {
        operation.input_schema.validate_shape(input)?;

        let mapped_input = operation.input_mapping.apply(&json!({ "mcp": input }))?;
        let prepared_request = PreparedRequest::from_mapping_output(&mapped_input)?;

        let adapter_response = match &operation.target {
            Target::Rest(target) => {
                let request = RestRequest {
                    path_params: prepared_request.path_params.clone(),
                    query_params: prepared_request.query_params.clone(),
                    headers: merge_headers(
                        &target.static_headers,
                        &operation.execution_config.headers,
                        &prepared_request.headers,
                    ),
                    body: prepared_request.body.clone(),
                    timeout_ms: operation.execution_config.timeout_ms,
                };
                let response = self.rest_adapter.execute(target, &request).await?;

                AdapterResponse {
                    status_code: response.status_code,
                    headers: response.headers,
                    body: response.body,
                }
            }
            _ => {
                return Err(RuntimeError::UnsupportedProtocol {
                    protocol: operation.protocol,
                });
            }
        };

        let finalized_output = finalize_output(operation, &adapter_response)?;

        operation.output_schema.validate_shape(&finalized_output)?;

        Ok(finalized_output)
    }
}

impl PreparedRequest {
    pub fn from_mapping_output(mapped: &Value) -> Result<Self, RuntimeError> {
        let request =
            mapped
                .get("request")
                .ok_or_else(|| RuntimeError::InvalidPreparedRequest {
                    details: "mapped input must contain request root".to_owned(),
                })?;

        Ok(Self {
            path_params: read_string_map(request.get("path"), "request.path")?,
            query_params: read_string_map(request.get("query"), "request.query")?,
            headers: read_string_map(request.get("headers"), "request.headers")?,
            body: request.get("body").and_then(non_empty_body).cloned(),
        })
    }
}

fn finalize_output(
    operation: &RuntimeOperation,
    response: &AdapterResponse,
) -> Result<Value, RuntimeError> {
    let mapped = operation.output_mapping.apply(&json!({
        "response": {
            "body": response.body,
            "data": response.body,
            "headers": response.headers,
            "status": response.status_code
        }
    }))?;

    Ok(mapped
        .get("output")
        .cloned()
        .unwrap_or_else(|| Value::Object(Map::new())))
}

fn merge_headers(
    static_headers: &BTreeMap<String, String>,
    execution_headers: &BTreeMap<String, String>,
    request_headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
    let mut headers = static_headers.clone();
    headers.extend(execution_headers.clone());
    headers.extend(request_headers.clone());
    headers
}

fn read_string_map(
    value: Option<&Value>,
    field_name: &str,
) -> Result<BTreeMap<String, String>, RuntimeError> {
    let Some(value) = value else {
        return Ok(BTreeMap::new());
    };

    let Some(object) = value.as_object() else {
        return Err(RuntimeError::InvalidPreparedRequest {
            details: format!("{field_name} must be an object"),
        });
    };

    object
        .iter()
        .map(|(key, value)| stringify_value(value).map(|value| (key.clone(), value)))
        .collect::<Result<BTreeMap<_, _>, _>>()
        .map_err(|details| RuntimeError::InvalidPreparedRequest { details })
}

fn stringify_value(value: &Value) -> Result<String, String> {
    match value {
        Value::Null => Ok("null".to_owned()),
        Value::Bool(value) => Ok(value.to_string()),
        Value::Number(value) => Ok(value.to_string()),
        Value::String(value) => Ok(value.clone()),
        Value::Array(_) | Value::Object(_) => {
            Err("request path/query/headers accept only scalar values".to_owned())
        }
    }
}

fn non_empty_body(value: &Value) -> Option<&Value> {
    match value {
        Value::Null => None,
        Value::Object(object) if object.is_empty() => None,
        _ => Some(value),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use axum::{Json, Router, routing::post};
    use mcpaas_core::{
        ExecutionConfig, GeneratedDraft, GeneratedDraftStatus, HttpMethod, Operation, OperationId,
        OperationStatus, Protocol, RestTarget, Samples, Target, ToolDescription, ToolExample,
    };
    use mcpaas_mapping::{MappingRule, MappingSet};
    use mcpaas_schema::{Schema, SchemaKind};
    use serde_json::{Value, json};
    use tokio::net::TcpListener;

    use crate::{RuntimeError, RuntimeExecutor, RuntimeOperation};

    #[tokio::test]
    async fn executes_rest_operation_end_to_end() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
        let operation = test_rest_operation(&base_url, false, false);

        let output = executor
            .execute(&operation, &json!({ "email": "user@example.com" }))
            .await
            .unwrap();

        assert_eq!(output, json!({ "id": "lead_123" }));
    }

    #[tokio::test]
    async fn rejects_invalid_input_shape() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
        let operation = test_rest_operation(&base_url, false, false);

        let error = executor.execute(&operation, &json!({})).await.unwrap_err();

        assert!(matches!(error, RuntimeError::Schema(_)));
    }

    #[tokio::test]
    async fn propagates_external_rest_errors() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
        let operation = test_rest_operation(&base_url, true, false);

        let error = executor
            .execute(&operation, &json!({ "email": "user@example.com" }))
            .await
            .unwrap_err();

        assert!(matches!(error, RuntimeError::RestAdapter(_)));
    }

    #[tokio::test]
    async fn fails_when_output_mapping_requires_missing_field() {
        let base_url = spawn_runtime_server().await;
        let executor = RuntimeExecutor::new();
        let operation = test_rest_operation(&base_url, false, true);

        let error = executor
            .execute(&operation, &json!({ "email": "user@example.com" }))
            .await
            .unwrap_err();

        assert!(matches!(error, RuntimeError::Mapping(_)));
    }

    async fn spawn_runtime_server() -> String {
        let app = Router::new().route("/leads", post(create_lead));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();

        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        format!("http://{}", address)
    }

    async fn create_lead(Json(payload): Json<Value>) -> (axum::http::StatusCode, Json<Value>) {
        let should_fail = payload
            .get("fail")
            .and_then(Value::as_bool)
            .unwrap_or(false);

        if should_fail {
            return (
                axum::http::StatusCode::BAD_GATEWAY,
                Json(json!({ "error": "upstream failed" })),
            );
        }

        (
            axum::http::StatusCode::OK,
            Json(json!({ "id": "lead_123", "email": payload["email"] })),
        )
    }

    fn test_rest_operation(
        base_url: &str,
        should_fail: bool,
        invalid_output: bool,
    ) -> RuntimeOperation {
        let mut input_rules = vec![MappingRule {
            source: "$.mcp.email".to_owned(),
            target: "$.request.body.email".to_owned(),
            required: true,
            default_value: None,
            transform: None,
            condition: None,
            notes: None,
        }];

        if should_fail {
            input_rules.push(MappingRule {
                source: "$.mcp.fail".to_owned(),
                target: "$.request.body.fail".to_owned(),
                required: false,
                default_value: Some(Value::Bool(true)),
                transform: None,
                condition: None,
                notes: None,
            });
        }

        let output_source = if invalid_output {
            "$.response.body.missing"
        } else {
            "$.response.body.id"
        };

        RuntimeOperation::from(Operation {
            id: OperationId::new("op_rest_runtime"),
            name: "crm_create_lead".to_owned(),
            display_name: "Create Lead".to_owned(),
            protocol: Protocol::Rest,
            status: OperationStatus::Published,
            version: 1,
            target: Target::Rest(RestTarget {
                base_url: base_url.to_owned(),
                method: HttpMethod::Post,
                path_template: "/leads".to_owned(),
                static_headers: BTreeMap::new(),
            }),
            input_schema: object_schema("email", SchemaKind::String),
            output_schema: object_schema("id", SchemaKind::String),
            input_mapping: MappingSet { rules: input_rules },
            output_mapping: MappingSet {
                rules: vec![MappingRule {
                    source: output_source.to_owned(),
                    target: "$.output.id".to_owned(),
                    required: true,
                    default_value: None,
                    transform: None,
                    condition: None,
                    notes: None,
                }],
            },
            execution_config: ExecutionConfig {
                timeout_ms: 1_000,
                retry_policy: None,
                auth_profile_ref: None,
                headers: BTreeMap::new(),
                protocol_options: None,
            },
            tool_description: ToolDescription {
                title: "Create Lead".to_owned(),
                description: "Creates a CRM lead".to_owned(),
                tags: vec!["crm".to_owned()],
                examples: vec![ToolExample {
                    input: json!({ "email": "user@example.com" }),
                }],
            },
            samples: Some(Samples::default()),
            generated_draft: Some(GeneratedDraft {
                status: GeneratedDraftStatus::Available,
                source_types: vec!["input_json".to_owned()],
                generated_at: Some("2026-03-25T20:00:00Z".to_owned()),
                input_schema_generated: true,
                output_schema_generated: true,
                input_mapping_generated: true,
                output_mapping_generated: true,
                warnings: Vec::new(),
            }),
            config_export: None,
            created_at: "2026-03-25T20:00:00Z".to_owned(),
            updated_at: "2026-03-25T20:00:00Z".to_owned(),
            published_at: Some("2026-03-25T20:00:00Z".to_owned()),
        })
    }

    fn object_schema(field_name: &str, kind: SchemaKind) -> Schema {
        Schema {
            kind: SchemaKind::Object,
            description: None,
            required: true,
            nullable: false,
            default_value: None,
            fields: BTreeMap::from([(
                field_name.to_owned(),
                Schema {
                    kind,
                    description: None,
                    required: true,
                    nullable: false,
                    default_value: None,
                    fields: BTreeMap::new(),
                    items: None,
                    enum_values: Vec::new(),
                    variants: Vec::new(),
                },
            )]),
            items: None,
            enum_values: Vec::new(),
            variants: Vec::new(),
        }
    }
}