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

use std::collections::BTreeMap;

use crank_core::{AgentId, InvocationSource, WorkspaceId};

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResponseCacheScope {
    pub workspace_key: String,
    pub agent_key: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RuntimeRequestContext {
    pub request_id: String,
    pub correlation_id: String,
    pub response_cache_scope: Option<ResponseCacheScope>,
    pub metering_context: Option<MeteringContext>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MeteringContext {
    pub workspace_id: WorkspaceId,
    pub agent_id: Option<AgentId>,
    pub source: InvocationSource,
}

impl RuntimeRequestContext {
    pub fn new(request_id: impl Into<String>, correlation_id: impl Into<String>) -> Self {
        Self {
            request_id: request_id.into(),
            correlation_id: correlation_id.into(),
            response_cache_scope: None,
            metering_context: None,
        }
    }

    pub fn from_request_id(request_id: impl Into<String>) -> Self {
        let request_id = request_id.into();
        Self::new(request_id.clone(), request_id)
    }

    pub fn outbound_headers(&self) -> BTreeMap<String, String> {
        BTreeMap::from([
            ("x-request-id".to_owned(), self.request_id.clone()),
            ("x-correlation-id".to_owned(), self.correlation_id.clone()),
        ])
    }

    pub fn with_response_cache_scope(
        mut self,
        workspace_key: impl Into<String>,
        agent_key: impl Into<String>,
    ) -> Self {
        self.response_cache_scope = Some(ResponseCacheScope {
            workspace_key: workspace_key.into(),
            agent_key: agent_key.into(),
        });
        self
    }

    pub fn response_cache_scope(&self) -> Option<&ResponseCacheScope> {
        self.response_cache_scope.as_ref()
    }

    pub fn with_metering_context(
        mut self,
        workspace_id: WorkspaceId,
        agent_id: Option<AgentId>,
        source: InvocationSource,
    ) -> Self {
        self.metering_context = Some(MeteringContext {
            workspace_id,
            agent_id,
            source,
        });
        self
    }

    pub fn metering_context(&self) -> Option<&MeteringContext> {
        self.metering_context.as_ref()
    }
}

impl From<ResponseCacheScope> for crank_core::ResponseCacheScope {
    fn from(value: ResponseCacheScope) -> Self {
        Self {
            workspace_key: value.workspace_key,
            agent_key: value.agent_key,
        }
    }
}

impl From<&ResponseCacheScope> for crank_core::ResponseCacheScope {
    fn from(value: &ResponseCacheScope) -> Self {
        Self {
            workspace_key: value.workspace_key.clone(),
            agent_key: value.agent_key.clone(),
        }
    }
}

impl From<&RuntimeRequestContext> for crank_core::RuntimeRequestContext {
    fn from(value: &RuntimeRequestContext) -> Self {
        Self {
            request_id: value.request_id.clone(),
            correlation_id: value.correlation_id.clone(),
            response_cache_scope: value.response_cache_scope.as_ref().map(Into::into),
            metering_context: value.metering_context.as_ref().map(Into::into),
        }
    }
}

impl From<MeteringContext> for crank_core::MeteringContext {
    fn from(value: MeteringContext) -> Self {
        Self {
            workspace_id: value.workspace_id,
            agent_id: value.agent_id,
            source: value.source,
        }
    }
}

impl From<&MeteringContext> for crank_core::MeteringContext {
    fn from(value: &MeteringContext) -> Self {
        Self {
            workspace_id: value.workspace_id.clone(),
            agent_id: value.agent_id.clone(),
            source: value.source,
        }
    }
}

#[cfg(test)]
mod tests {
    use crank_core::{InvocationSource, WorkspaceId};

    use super::RuntimeRequestContext;

    #[test]
    fn uses_request_id_for_default_correlation_id() {
        let context = RuntimeRequestContext::from_request_id("req_123");

        assert_eq!(context.request_id, "req_123");
        assert_eq!(context.correlation_id, "req_123");
        assert!(context.response_cache_scope.is_none());
    }

    #[test]
    fn attaches_response_cache_scope() {
        let context = RuntimeRequestContext::from_request_id("req_123")
            .with_response_cache_scope("ws_01", "agent_01");

        let scope = context.response_cache_scope().unwrap();
        assert_eq!(scope.workspace_key, "ws_01");
        assert_eq!(scope.agent_key, "agent_01");
    }

    #[test]
    fn attaches_metering_context() {
        let context = RuntimeRequestContext::from_request_id("req_123").with_metering_context(
            WorkspaceId::new("ws_01"),
            None,
            InvocationSource::AdminTestRun,
        );

        let metering = context.metering_context().unwrap();
        assert_eq!(metering.workspace_id, WorkspaceId::new("ws_01"));
        assert_eq!(metering.agent_id, None);
        assert_eq!(metering.source, InvocationSource::AdminTestRun);
    }
}