crank/crates/crank-runtime/src/execution_request.rs

use std::time::Instant;

use crank_core::{AgentId, ExecutionOrigin, InvocationSource, WorkspaceId};
use serde_json::Value;

use crate::{ResolvedAuth, RuntimeOperation, RuntimeRequestContext};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExecutionAuthorization {
    Authorized,
}

#[derive(Clone, Copy)]
pub struct RuntimeExecutionRequest<'a> {
    workspace_id: &'a WorkspaceId,
    origin: ExecutionOrigin,
    agent_id: Option<&'a AgentId>,
    operation: &'a RuntimeOperation,
    input: &'a Value,
    authorization: ExecutionAuthorization,
    resolved_auth: Option<&'a ResolvedAuth>,
    request_context: &'a RuntimeRequestContext,
    deadline: Instant,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ExecutionRequestError {
    #[error("execution origin and agent identity are incompatible")]
    InvalidOrigin,
    #[error("operation version must be positive")]
    InvalidOperationVersion,
    #[error("execution deadline has already elapsed")]
    DeadlineElapsed,
    #[error("execution metering scope is missing or incompatible")]
    InvalidMeteringScope,
}

impl<'a> RuntimeExecutionRequest<'a> {
    #[allow(clippy::too_many_arguments)]
    pub fn try_new(
        workspace_id: &'a WorkspaceId,
        origin: ExecutionOrigin,
        agent_id: Option<&'a AgentId>,
        operation: &'a RuntimeOperation,
        input: &'a Value,
        authorization: ExecutionAuthorization,
        resolved_auth: Option<&'a ResolvedAuth>,
        request_context: &'a RuntimeRequestContext,
        deadline: Instant,
    ) -> Result<Self, ExecutionRequestError> {
        origin
            .validate_agent(agent_id)
            .map_err(|_| ExecutionRequestError::InvalidOrigin)?;
        if operation.operation_version == 0 {
            return Err(ExecutionRequestError::InvalidOperationVersion);
        }
        if deadline <= Instant::now() {
            return Err(ExecutionRequestError::DeadlineElapsed);
        }
        let expected_source = match origin {
            ExecutionOrigin::AdminDraft => InvocationSource::AdminTestRun,
            ExecutionOrigin::AgentSnapshot => InvocationSource::AgentToolCall,
        };
        let metering = request_context
            .metering_context()
            .ok_or(ExecutionRequestError::InvalidMeteringScope)?;
        if &metering.workspace_id != workspace_id
            || metering.source != expected_source
            || metering.agent_id.as_ref() != agent_id
        {
            return Err(ExecutionRequestError::InvalidMeteringScope);
        }

        Ok(Self {
            workspace_id,
            origin,
            agent_id,
            operation,
            input,
            authorization,
            resolved_auth,
            request_context,
            deadline,
        })
    }

    pub(crate) fn operation(&self) -> &'a RuntimeOperation {
        self.operation
    }

    pub(crate) fn workspace_id(&self) -> &'a WorkspaceId {
        self.workspace_id
    }

    pub(crate) fn agent_id(&self) -> Option<&'a AgentId> {
        self.agent_id
    }

    pub(crate) fn authorization(&self) -> ExecutionAuthorization {
        self.authorization
    }

    pub(crate) fn input(&self) -> &'a Value {
        self.input
    }

    pub(crate) fn origin(&self) -> ExecutionOrigin {
        self.origin
    }

    pub(crate) fn resolved_auth(&self) -> Option<&'a ResolvedAuth> {
        self.resolved_auth
    }

    pub(crate) fn request_context(&self) -> &'a RuntimeRequestContext {
        self.request_context
    }

    pub(crate) fn deadline(&self) -> Instant {
        self.deadline
    }
}

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

    use crank_core::{
        AgentId, ExecutionConfig, HttpMethod, OperationId, RestTarget, Target, ToolDescription,
    };
    use crank_mapping::MappingSet;
    use crank_schema::{Schema, SchemaKind};
    use serde_json::json;

    use super::*;

    #[test]
    fn rejects_origin_or_scope_mismatch_before_execution() {
        let operation = operation();
        let input = json!({});
        let workspace = WorkspaceId::new("ws_01");
        let foreign = WorkspaceId::new("ws_02");
        let agent = AgentId::new("agent_01");
        let context = RuntimeRequestContext::from_request_id("req_01").with_metering_context(
            workspace.clone(),
            None,
            InvocationSource::AdminTestRun,
        );
        let deadline = Instant::now() + Duration::from_secs(1);

        assert_eq!(
            RuntimeExecutionRequest::try_new(
                &workspace,
                ExecutionOrigin::AdminDraft,
                Some(&agent),
                &operation,
                &input,
                ExecutionAuthorization::Authorized,
                None,
                &context,
                deadline,
            )
            .err()
            .expect("origin mismatch"),
            ExecutionRequestError::InvalidOrigin
        );
        assert_eq!(
            RuntimeExecutionRequest::try_new(
                &foreign,
                ExecutionOrigin::AdminDraft,
                None,
                &operation,
                &input,
                ExecutionAuthorization::Authorized,
                None,
                &context,
                deadline,
            )
            .err()
            .expect("scope mismatch"),
            ExecutionRequestError::InvalidMeteringScope
        );
    }

    fn operation() -> RuntimeOperation {
        let schema = Schema {
            kind: SchemaKind::Object,
            description: None,
            required: true,
            nullable: false,
            default_value: None,
            fields: BTreeMap::new(),
            items: None,
            enum_values: Vec::new(),
            variants: Vec::new(),
        };
        RuntimeOperation {
            operation_id: OperationId::new("op_01"),
            operation_version: 1,
            tool_name: "tool".to_owned(),
            protocol: crank_core::Protocol::Rest,
            target: Target::Rest(RestTarget {
                base_url: "https://example.invalid".to_owned(),
                method: HttpMethod::Post,
                path_template: "/".to_owned(),
                static_headers: BTreeMap::new(),
            }),
            input_schema: schema.clone(),
            output_schema: schema,
            input_mapping: MappingSet::default(),
            output_mapping: MappingSet::default(),
            execution_config: ExecutionConfig {
                timeout_ms: 1_000,
                retry_policy: None,
                response_cache: None,
                idempotency: None,
                safety: None,
                approval_policy: None,
                auth_profile_ref: None,
                headers: BTreeMap::new(),
            },
            tool_description: ToolDescription {
                title: "Tool".to_owned(),
                description: "Tool".to_owned(),
                tags: Vec::new(),
                examples: Vec::new(),
            },
        }
    }
}