runtime: emit metering events via public seam
This commit is contained in:
@@ -168,5 +168,6 @@ Progress:
|
||||
- Phase 3 dependency slice: added public tenancy seam in `crank-core` — `TenantId`, `TenantResolutionContext`, `TenantController`, `TenancyError`, and `SingleTenantController` — so `cloud` can build hosted tenant routing without bypassing the shared extension model
|
||||
- Phase 3 dependency slice: added public metering seam in `crank-core` — `MeteringEvent`, `MeteringSink`, `SharedMeteringSink`, and `NoopMeteringSink` — so `cloud` can add hosted usage capture without introducing private-only contracts
|
||||
- Phase 3 dependency slice: added public billing seam in `crank-core` — `BillingHook`, `BillingGate`, `BillingError`, `SharedBillingHook`, and `NoopBillingHook` — so `cloud` can layer billing policy without private-only contracts in the base repo
|
||||
- Phase 3 runtime slice: `RuntimeRequestContext` now carries optional metering context (`workspace_id`, `agent_id`, `source`), `RuntimeExecutorBuilder` accepts a `MeteringSink`, and `RuntimeExecutor` emits `MeteringEvent` on invocation completion while Community apps keep using the default `NoopMeteringSink`
|
||||
- Phase 3 dependency slice: added `CacheBackendFactory` in `crank-runtime` with `BuiltinCacheBackendFactory`; unlike the original draft, the seam lives in runtime rather than core to avoid a `crank-core -> crank-runtime` dependency cycle around `RuntimeCacheStores`
|
||||
- backward-compatible alias `CommunityMachineCredentialVerifier` сохранен, поведение Community не изменено
|
||||
|
||||
@@ -2023,7 +2023,12 @@ impl AdminService {
|
||||
payload: TestRunPayload,
|
||||
request_id: &str,
|
||||
) -> Result<TestRunResult, ApiError> {
|
||||
let runtime_request_context = RuntimeRequestContext::from_request_id(request_id);
|
||||
let runtime_request_context = RuntimeRequestContext::from_request_id(request_id)
|
||||
.with_metering_context(
|
||||
workspace_id.clone(),
|
||||
None,
|
||||
InvocationSource::AdminTestRun,
|
||||
);
|
||||
let record = self
|
||||
.get_operation_version(workspace_id, operation_id, payload.version)
|
||||
.await?;
|
||||
|
||||
@@ -710,6 +710,11 @@ async fn handle_base_tool_call(
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
);
|
||||
let request_preview = build_request_preview(&state.runtime, &operation, &arguments);
|
||||
let started_at = Instant::now();
|
||||
@@ -825,6 +830,11 @@ async fn handle_session_start_call(
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
);
|
||||
let request_preview = build_request_preview(&state.runtime, &runtime_operation, &arguments);
|
||||
let started_at = Instant::now();
|
||||
@@ -1256,6 +1266,11 @@ async fn handle_async_job_start_call(
|
||||
.with_response_cache_scope(
|
||||
tool.workspace_id.as_str().to_owned(),
|
||||
tool.agent_id.as_str().to_owned(),
|
||||
)
|
||||
.with_metering_context(
|
||||
tool.workspace_id.clone(),
|
||||
Some(tool.agent_id.clone()),
|
||||
InvocationSource::AgentToolCall,
|
||||
);
|
||||
let request_preview = build_request_preview(&state.runtime, &operation, &arguments);
|
||||
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
|
||||
|
||||
@@ -4,7 +4,7 @@ use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{ExecutionMode, Protocol, StreamSession, Target};
|
||||
use crate::{AgentId, ExecutionMode, InvocationSource, Protocol, StreamSession, Target, WorkspaceId};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Default)]
|
||||
pub struct ResponseCacheScope {
|
||||
@@ -17,6 +17,14 @@ 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 {
|
||||
@@ -25,6 +33,7 @@ impl RuntimeRequestContext {
|
||||
request_id: request_id.into(),
|
||||
correlation_id: correlation_id.into(),
|
||||
response_cache_scope: None,
|
||||
metering_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -55,6 +64,24 @@ impl RuntimeRequestContext {
|
||||
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)]
|
||||
@@ -191,7 +218,9 @@ mod tests {
|
||||
AdapterRegistry, AdapterResponse, PreparedRequest, ProtocolAdapter, ProtocolAdapterError,
|
||||
RuntimeRequestContext,
|
||||
};
|
||||
use crate::{ExecutionMode, HttpMethod, Protocol, RestTarget, Target};
|
||||
use crate::{
|
||||
ExecutionMode, HttpMethod, InvocationSource, Protocol, RestTarget, Target, WorkspaceId,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StubAdapter(Protocol);
|
||||
@@ -253,6 +282,20 @@ mod tests {
|
||||
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)));
|
||||
|
||||
@@ -54,7 +54,8 @@ pub use ext::capability::{CapabilityProfile, CommunityCapabilityProfile};
|
||||
pub use ext::metering::{MeteringEvent, MeteringSink, NoopMeteringSink, SharedMeteringSink};
|
||||
pub use ext::protocol::{
|
||||
AdapterRegistry, AdapterResponse, PreparedRequest, ProtocolAdapter, ProtocolAdapterError,
|
||||
ResponseCacheScope, RuntimeRequestContext, SharedProtocolAdapter, WindowExecutionResult,
|
||||
MeteringContext, ResponseCacheScope, RuntimeRequestContext, SharedProtocolAdapter,
|
||||
WindowExecutionResult,
|
||||
};
|
||||
pub use ext::tenancy::{
|
||||
SharedTenantController, SingleTenantController, TenantController, TenantResolutionContext,
|
||||
|
||||
@@ -25,6 +25,7 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
time.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
tracing.workspace = true
|
||||
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use crank_core::{
|
||||
AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, ResponseCacheStore,
|
||||
SharedProtocolAdapter, Target, TransportBehavior,
|
||||
AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, InvocationStatus,
|
||||
MeteringEvent, ResponseCacheStore, SharedMeteringSink, SharedProtocolAdapter, Target,
|
||||
TransportBehavior,
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use time::OffsetDateTime;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
use tracing::debug;
|
||||
|
||||
@@ -26,6 +28,7 @@ pub struct RuntimeExecutor {
|
||||
session_limiter: Arc<Semaphore>,
|
||||
job_limiter: Arc<Semaphore>,
|
||||
response_cache: Option<Arc<dyn ResponseCacheStore>>,
|
||||
metering_sink: SharedMeteringSink,
|
||||
}
|
||||
|
||||
impl Default for RuntimeExecutor {
|
||||
@@ -49,6 +52,7 @@ impl RuntimeExecutor {
|
||||
limits: RuntimeLimits,
|
||||
adapters: AdapterRegistry,
|
||||
response_cache: Option<Arc<dyn ResponseCacheStore>>,
|
||||
metering_sink: SharedMeteringSink,
|
||||
) -> Self {
|
||||
Self {
|
||||
adapters,
|
||||
@@ -58,6 +62,7 @@ impl RuntimeExecutor {
|
||||
job_limiter: Arc::new(Semaphore::new(limits.max_concurrent_jobs)),
|
||||
limits,
|
||||
response_cache,
|
||||
metering_sink,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -107,10 +112,15 @@ impl RuntimeExecutor {
|
||||
) -> Result<Value, RuntimeError> {
|
||||
log_runtime_event("unary.execute", operation, request_context);
|
||||
let _permit = self.acquire_unary_permit(operation)?;
|
||||
let started_at = Instant::now();
|
||||
let prepared_request = self.prepare_request(operation, input)?;
|
||||
let prepared_request = apply_resolved_auth(prepared_request, resolved_auth);
|
||||
self.execute_prepared(operation, prepared_request, request_context)
|
||||
.await
|
||||
let result = self
|
||||
.execute_prepared(operation, prepared_request, request_context)
|
||||
.await;
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn execute_window(
|
||||
@@ -151,17 +161,24 @@ impl RuntimeExecutor {
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
log_runtime_event("window.execute", operation, request_context);
|
||||
let _permit = self.acquire_window_permit()?;
|
||||
let started_at = Instant::now();
|
||||
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
|
||||
return Err(RuntimeError::MissingStreamingConfig {
|
||||
let result = Err(RuntimeError::MissingStreamingConfig {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
});
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
return result;
|
||||
};
|
||||
|
||||
if streaming.mode != ExecutionMode::Window {
|
||||
return Err(RuntimeError::UnsupportedExecutionMode {
|
||||
let result = Err(RuntimeError::UnsupportedExecutionMode {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
mode: streaming.mode,
|
||||
});
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
return result;
|
||||
}
|
||||
|
||||
let prepared_request = self.prepare_request(operation, input)?;
|
||||
@@ -177,7 +194,10 @@ impl RuntimeExecutor {
|
||||
.await?
|
||||
};
|
||||
|
||||
crate::aggregation::collect_window_result(&adapter_response.body, streaming)
|
||||
let result = crate::aggregation::collect_window_result(&adapter_response.body, streaming);
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn execute_session_seed(
|
||||
@@ -296,6 +316,34 @@ impl RuntimeExecutor {
|
||||
Ok(finalized_output)
|
||||
}
|
||||
|
||||
async fn record_metering<T>(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
result: &Result<T, RuntimeError>,
|
||||
started_at: Instant,
|
||||
) {
|
||||
let Some(context) = request_context.and_then(RuntimeRequestContext::metering_context) else {
|
||||
return;
|
||||
};
|
||||
|
||||
self.metering_sink
|
||||
.record(MeteringEvent {
|
||||
workspace_id: context.workspace_id.clone(),
|
||||
agent_id: context.agent_id.clone(),
|
||||
operation_id: operation.operation_id.clone(),
|
||||
source: context.source,
|
||||
status: if result.is_ok() {
|
||||
InvocationStatus::Ok
|
||||
} else {
|
||||
InvocationStatus::Error
|
||||
},
|
||||
duration_ms: started_at.elapsed().as_millis() as u64,
|
||||
created_at: OffsetDateTime::now_utc(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn execute_adapter(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
@@ -812,6 +860,7 @@ mod tests {
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::Query,
|
||||
@@ -821,9 +870,10 @@ mod tests {
|
||||
#[cfg(any())]
|
||||
use crank_adapter_grpc::test_support as grpc_test_support;
|
||||
use crank_core::{
|
||||
ExecutionConfig, GeneratedDraft, GeneratedDraftStatus, HttpMethod, Operation, OperationId,
|
||||
ExecutionConfig, GeneratedDraft, GeneratedDraftStatus, HttpMethod, InvocationSource,
|
||||
InvocationStatus, MeteringEvent, MeteringSink, Operation, OperationId,
|
||||
OperationSecurityLevel, OperationStatus, Protocol, RestTarget, Samples, Target,
|
||||
ToolDescription, ToolExample,
|
||||
ToolDescription, ToolExample, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::{MappingRule, MappingSet};
|
||||
use crank_schema::{Schema, SchemaKind};
|
||||
@@ -838,8 +888,8 @@ mod tests {
|
||||
use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*};
|
||||
|
||||
use crate::{
|
||||
InMemoryResponseCacheStore, PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeLimits,
|
||||
RuntimeOperation, RuntimeRequestContext,
|
||||
InMemoryResponseCacheStore, PreparedRequest, RuntimeError, RuntimeExecutor,
|
||||
RuntimeExecutorBuilder, RuntimeLimits, RuntimeOperation, RuntimeRequestContext,
|
||||
};
|
||||
|
||||
fn timestamp(value: &str) -> OffsetDateTime {
|
||||
@@ -882,6 +932,24 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct RecordingMeteringSink {
|
||||
events: Arc<Mutex<Vec<MeteringEvent>>>,
|
||||
}
|
||||
|
||||
impl RecordingMeteringSink {
|
||||
fn events(&self) -> Vec<MeteringEvent> {
|
||||
self.events.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl MeteringSink for RecordingMeteringSink {
|
||||
async fn record(&self, event: MeteringEvent) {
|
||||
self.events.lock().unwrap().push(event);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executes_rest_operation_end_to_end() {
|
||||
let base_url = spawn_runtime_server().await;
|
||||
@@ -919,6 +987,42 @@ mod tests {
|
||||
assert_eq!(output, json!({ "id": "req_rest_123|corr_rest_123" }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn records_metering_event_for_successful_invocation() {
|
||||
let base_url = spawn_runtime_server().await;
|
||||
let metering_sink = RecordingMeteringSink::default();
|
||||
let executor = RuntimeExecutorBuilder::new()
|
||||
.register_adapter(Arc::new(crank_adapter_rest::RestAdapter::new()))
|
||||
.with_metering_sink(Arc::new(metering_sink.clone()))
|
||||
.build();
|
||||
let operation = test_rest_operation(&base_url, false, false);
|
||||
let context = RuntimeRequestContext::from_request_id("req_meter_1").with_metering_context(
|
||||
WorkspaceId::new("ws_meter"),
|
||||
None,
|
||||
InvocationSource::AgentToolCall,
|
||||
);
|
||||
|
||||
let output = executor
|
||||
.execute_with_context(
|
||||
&operation,
|
||||
&json!({ "email": "user@example.com" }),
|
||||
Some(&context),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(output, json!({ "id": "lead_123" }));
|
||||
|
||||
let events = metering_sink.events();
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(events[0].workspace_id, WorkspaceId::new("ws_meter"));
|
||||
assert_eq!(events[0].agent_id, None);
|
||||
assert_eq!(events[0].operation_id, operation.operation_id);
|
||||
assert_eq!(events[0].source, InvocationSource::AgentToolCall);
|
||||
assert_eq!(events[0].status, InvocationStatus::Ok);
|
||||
assert!(events[0].duration_ms <= 60_000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn caches_rest_get_responses_within_agent_scope() {
|
||||
let request_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crank_adapter_rest::RestAdapter;
|
||||
use crank_core::{AdapterRegistry, ResponseCacheStore, SharedProtocolAdapter};
|
||||
use crank_core::{
|
||||
AdapterRegistry, MeteringSink, NoopMeteringSink, ResponseCacheStore, SharedMeteringSink,
|
||||
SharedProtocolAdapter,
|
||||
};
|
||||
|
||||
use crate::{RuntimeExecutor, RuntimeLimits};
|
||||
|
||||
@@ -9,6 +12,7 @@ pub struct RuntimeExecutorBuilder {
|
||||
limits: RuntimeLimits,
|
||||
adapters: AdapterRegistry,
|
||||
response_cache: Option<Arc<dyn ResponseCacheStore>>,
|
||||
metering_sink: SharedMeteringSink,
|
||||
}
|
||||
|
||||
impl RuntimeExecutorBuilder {
|
||||
@@ -17,6 +21,7 @@ impl RuntimeExecutorBuilder {
|
||||
limits: RuntimeLimits::default(),
|
||||
adapters: AdapterRegistry::empty(),
|
||||
response_cache: None,
|
||||
metering_sink: Arc::new(NoopMeteringSink) as Arc<dyn MeteringSink>,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,8 +40,18 @@ impl RuntimeExecutorBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_metering_sink(mut self, sink: SharedMeteringSink) -> Self {
|
||||
self.metering_sink = sink;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> RuntimeExecutor {
|
||||
RuntimeExecutor::from_builder_parts(self.limits, self.adapters, self.response_cache)
|
||||
RuntimeExecutor::from_builder_parts(
|
||||
self.limits,
|
||||
self.adapters,
|
||||
self.response_cache,
|
||||
self.metering_sink,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -30,6 +30,6 @@ pub use model::{AdapterResponse, PreparedRequest, RuntimeOperation};
|
||||
pub use rate_limit::{
|
||||
RateLimitRejection, RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter,
|
||||
};
|
||||
pub use request_context::{ResponseCacheScope, RuntimeRequestContext};
|
||||
pub use request_context::{MeteringContext, ResponseCacheScope, RuntimeRequestContext};
|
||||
pub use secret_crypto::SecretCrypto;
|
||||
pub use streaming::WindowExecutionResult;
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crank_core::{AgentId, InvocationSource, WorkspaceId};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ResponseCacheScope {
|
||||
pub workspace_key: String,
|
||||
@@ -11,6 +13,14 @@ 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 {
|
||||
@@ -19,6 +29,7 @@ impl RuntimeRequestContext {
|
||||
request_id: request_id.into(),
|
||||
correlation_id: correlation_id.into(),
|
||||
response_cache_scope: None,
|
||||
metering_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,6 +60,24 @@ impl RuntimeRequestContext {
|
||||
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 {
|
||||
@@ -75,12 +104,35 @@ impl From<&RuntimeRequestContext> for crank_core::RuntimeRequestContext {
|
||||
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]
|
||||
@@ -101,4 +153,18 @@ mod tests {
|
||||
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