crank/crates/crank-core/src/ext/protocol.rs

use std::{collections::BTreeMap, sync::Arc};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::{
    AgentId, ExecutionMode, InvocationSource, Protocol, StreamSession, Target, WorkspaceId,
};

#[derive(Clone, Debug, PartialEq, Eq, Default)]
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()
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
pub struct PreparedRequest {
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub path_params: BTreeMap<String, String>,
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub query_params: BTreeMap<String, String>,
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub headers: BTreeMap<String, String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub grpc: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub variables: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub body: Option<Value>,
    #[serde(default)]
    pub timeout_ms: u64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AdapterResponse {
    pub status_code: u16,
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub headers: BTreeMap<String, String>,
    pub body: Value,
    pub data: Value,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WindowExecutionResult {
    pub summary: Value,
    pub items: Vec<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cursor: Option<Value>,
    pub window_complete: bool,
    pub truncated: bool,
    pub has_more: bool,
}

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ProtocolAdapterError {
    #[error("protocol {protocol:?} does not support execution mode {mode:?}")]
    UnsupportedMode {
        protocol: Protocol,
        mode: ExecutionMode,
    },
    #[error("{0}")]
    Message(String),
}

#[async_trait]
pub trait ProtocolAdapter: Send + Sync {
    fn protocol(&self) -> Protocol;

    fn supports_mode(&self, mode: ExecutionMode) -> bool;

    async fn invoke_unary(
        &self,
        target: &Target,
        prepared: &PreparedRequest,
        context: &RuntimeRequestContext,
    ) -> Result<AdapterResponse, ProtocolAdapterError>;

    async fn invoke_window(
        &self,
        _target: &Target,
        _prepared: &PreparedRequest,
        _window_duration_ms: u64,
        _max_items: Option<u32>,
        _context: &RuntimeRequestContext,
    ) -> Result<WindowExecutionResult, ProtocolAdapterError> {
        Err(ProtocolAdapterError::UnsupportedMode {
            protocol: self.protocol(),
            mode: ExecutionMode::Window,
        })
    }

    async fn open_session(
        &self,
        _target: &Target,
        _prepared: &PreparedRequest,
        _context: &RuntimeRequestContext,
    ) -> Result<StreamSession, ProtocolAdapterError> {
        Err(ProtocolAdapterError::UnsupportedMode {
            protocol: self.protocol(),
            mode: ExecutionMode::Session,
        })
    }
}

pub type SharedProtocolAdapter = Arc<dyn ProtocolAdapter>;

#[derive(Clone, Default)]
pub struct AdapterRegistry {
    adapters: Vec<SharedProtocolAdapter>,
}

impl AdapterRegistry {
    pub fn empty() -> Self {
        Self::default()
    }

    pub fn register(mut self, adapter: SharedProtocolAdapter) -> Self {
        self.adapters
            .retain(|existing| existing.protocol() != adapter.protocol());
        self.adapters.push(adapter);
        self
    }

    pub fn get(&self, protocol: Protocol) -> Option<SharedProtocolAdapter> {
        self.adapters
            .iter()
            .find(|adapter| adapter.protocol() == protocol)
            .cloned()
    }

    pub fn protocols(&self) -> Vec<Protocol> {
        self.adapters
            .iter()
            .map(|adapter| adapter.protocol())
            .collect()
    }
}

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

    use async_trait::async_trait;
    use serde_json::json;

    use super::{
        AdapterRegistry, AdapterResponse, PreparedRequest, ProtocolAdapter, ProtocolAdapterError,
        RuntimeRequestContext,
    };
    use crate::{
        ExecutionMode, HttpMethod, InvocationSource, Protocol, RestTarget, Target, WorkspaceId,
    };

    #[derive(Clone)]
    struct StubAdapter(Protocol);

    #[async_trait]
    impl ProtocolAdapter for StubAdapter {
        fn protocol(&self) -> Protocol {
            self.0
        }

        fn supports_mode(&self, mode: ExecutionMode) -> bool {
            matches!(mode, ExecutionMode::Unary)
        }

        async fn invoke_unary(
            &self,
            _target: &Target,
            _prepared: &PreparedRequest,
            _context: &RuntimeRequestContext,
        ) -> Result<AdapterResponse, ProtocolAdapterError> {
            Ok(AdapterResponse {
                status_code: 200,
                headers: BTreeMap::new(),
                body: json!({"ok": true}),
                data: json!({"ok": true}),
            })
        }
    }

    #[test]
    fn registry_returns_registered_protocols() {
        let registry = AdapterRegistry::empty()
            .register(Arc::new(StubAdapter(Protocol::Rest)))
            .register(Arc::new(StubAdapter(Protocol::Graphql)));

        assert_eq!(
            registry.protocols(),
            vec![Protocol::Rest, Protocol::Graphql]
        );
    }

    #[test]
    fn registry_replaces_adapter_for_same_protocol() {
        let registry = AdapterRegistry::empty()
            .register(Arc::new(StubAdapter(Protocol::Rest)))
            .register(Arc::new(StubAdapter(Protocol::Rest)));

        assert_eq!(registry.protocols(), vec![Protocol::Rest]);
        assert!(registry.get(Protocol::Rest).is_some());
    }

    #[test]
    fn request_context_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 request_context_attaches_metering_context() {
        let context = RuntimeRequestContext::from_request_id("req_123").with_metering_context(
            WorkspaceId::new("ws_01"),
            None,
            InvocationSource::AgentToolCall,
        );

        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::AgentToolCall);
    }

    #[tokio::test]
    async fn registry_returns_registered_adapter() {
        let registry = AdapterRegistry::empty().register(Arc::new(StubAdapter(Protocol::Rest)));
        let adapter = registry.get(Protocol::Rest).unwrap();

        let response = adapter
            .invoke_unary(
                &Target::Rest(RestTarget {
                    base_url: "https://example.test".to_owned(),
                    method: HttpMethod::Get,
                    path_template: "/health".to_owned(),
                    static_headers: BTreeMap::new(),
                }),
                &PreparedRequest::default(),
                &RuntimeRequestContext::from_request_id("req_1"),
            )
            .await
            .unwrap();

        assert_eq!(response.status_code, 200);
    }
}