chore: publish clean community baseline
This commit is contained in:
@@ -0,0 +1,170 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user