use serde_json::Value;
fn is_embedding_bearing(verb: &str, params: &Value) -> bool {
match verb {
"create" => params.get("items").is_none(),
"update" | "memory.remember" | "memory.recall" | "knowledge.search"
| "knowledge.compose" | "knowledge.index" => true,
_ => false,
}
}
fn base_weight(_verb: &str) -> i64 {
1
}
fn per_item_weight(verb: &str, params: &Value) -> i64 {
if is_embedding_bearing(verb, params) {
1
} else {
0
}
}
fn item_count(verb: &str, result: &Value) -> i64 {
if verb == "knowledge.index" {
result.get("total").and_then(Value::as_i64).unwrap_or(0)
} else {
1
}
}
fn model_count(verb: &str, params: &Value, registered_model_count: impl FnOnce() -> i64) -> i64 {
match verb {
"memory.remember" => {
let explicit_single_model = params
.get("embedding_model")
.and_then(Value::as_str)
.is_some();
if explicit_single_model {
1
} else {
registered_model_count()
}
}
"create" => registered_model_count(),
_ => 1,
}
}
fn compute(base_weight: i64, per_item_weight: i64, item_count: i64, model_count: i64) -> i64 {
let term = per_item_weight
.checked_mul(item_count)
.and_then(|p| p.checked_mul(model_count))
.unwrap_or(i64::MAX);
base_weight.checked_add(term).unwrap_or(i64::MAX)
}
pub fn cost_unit_for_dispatch(
verb: &str,
params: &Value,
result: &Value,
registered_model_count: impl FnOnce() -> i64,
) -> i64 {
let weight = per_item_weight(verb, params);
if weight == 0 {
return compute(base_weight(verb), 0, 0, 0);
}
let items = item_count(verb, result);
let models = model_count(verb, params, registered_model_count);
compute(base_weight(verb), weight, items, models)
}
pub fn resource_payload(
verb: &str,
params: &Value,
result: &Value,
registered_model_count: impl FnOnce() -> i64,
request_id: Option<u64>,
) -> Value {
let cost_unit = cost_unit_for_dispatch(verb, params, result, registered_model_count);
let mut payload = serde_json::json!({ "work_class": "interactive", "cost_unit": cost_unit });
if let Some(id) = request_id {
if let Value::Object(ref mut map) = payload {
map.insert("request_id".to_string(), serde_json::json!(id));
}
}
payload
}
pub fn base_resource_payload(request_id: Option<u64>) -> Value {
let mut payload = serde_json::json!({ "work_class": "interactive" });
if let Some(id) = request_id {
if let Value::Object(ref mut map) = payload {
map.insert("request_id".to_string(), serde_json::json!(id));
}
}
payload
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn unreachable_model_count() -> i64 {
panic!("registered_model_count must not be called for a non-embedding-bearing verb")
}
#[test]
fn non_embedding_verb_is_base_weight_only() {
let cost = cost_unit_for_dispatch("stats", &json!({}), &json!({}), unreachable_model_count);
assert_eq!(
cost, 1,
"non-embedding-bearing verb must be base_weight(verb) alone"
);
}
#[test]
fn link_is_base_weight_only_regardless_of_bulk_shape() {
let singleton = cost_unit_for_dispatch(
"link",
&json!({"source_id": "a", "target_id": "b", "relation": "extends"}),
&json!({"id": "edge-1"}),
unreachable_model_count,
);
let bulk = cost_unit_for_dispatch(
"link",
&json!({"links": [{}, {}, {}]}),
&json!({"attempted": 3, "created": 3}),
unreachable_model_count,
);
assert_eq!(singleton, 1);
assert_eq!(
bulk, 1,
"link has no embedding-bearing path, singleton or bulk"
);
}
#[test]
fn create_singleton_scales_with_registered_model_count() {
let cost = cost_unit_for_dispatch("create", &json!({"kind": "concept"}), &json!({}), || 3);
assert_eq!(cost, 4);
}
#[test]
fn zero_registered_models_vanishes_the_term() {
let cost = cost_unit_for_dispatch("create", &json!({"kind": "concept"}), &json!({}), || 0);
assert_eq!(
cost, 1,
"no embedding model registered -> cost_unit reduces to base_weight(verb) alone"
);
}
#[test]
fn memory_remember_zero_registered_models_vanishes_the_term() {
let cost = cost_unit_for_dispatch("memory.remember", &json!({}), &json!({}), || 0);
assert_eq!(cost, 1);
}
#[test]
fn bulk_create_is_base_weight_only_never_touches_model_count() {
let cost = cost_unit_for_dispatch(
"create",
&json!({"items": [{"kind": "concept", "name": "a"}, {"kind": "concept", "name": "b"}]}),
&json!({"attempted": 2, "created": 2}),
unreachable_model_count,
);
assert_eq!(
cost, 1,
"bulk create(items=[...]) skips embedding entirely -> base_weight(verb) alone"
);
}
#[test]
fn bulk_create_is_base_only_regardless_of_created_count() {
let items: Vec<Value> = (0..250)
.map(|i| json!({"kind": "concept", "name": format!("item-{i}")}))
.collect();
let cost = cost_unit_for_dispatch(
"create",
&json!({"items": items}),
&json!({"attempted": 250, "created": 250}),
unreachable_model_count,
);
assert_eq!(cost, 1);
}
#[test]
fn knowledge_index_uses_full_paged_total_not_batch_size_ceiling() {
let cost = cost_unit_for_dispatch(
"knowledge.index",
&json!({"batch_size": 1000}),
&json!({"indexed": 4500, "skipped": 0, "failed": 0, "total": 4500}),
|| 1,
);
assert_eq!(cost, 4501);
}
#[test]
fn knowledge_index_missing_total_defaults_to_zero_items_not_a_panic() {
let cost = cost_unit_for_dispatch("knowledge.index", &json!({}), &json!({}), || 1);
assert_eq!(cost, 1);
}
#[test]
fn knowledge_index_model_count_is_constant_one_never_reads_registry() {
let cost = cost_unit_for_dispatch(
"knowledge.index",
&json!({}),
&json!({"total": 10}),
unreachable_model_count,
);
assert_eq!(cost, 11);
}
#[test]
fn memory_remember_explicit_model_overrides_registry_count() {
let cost = cost_unit_for_dispatch(
"memory.remember",
&json!({"content": "x", "embedding_model": "paraphrase"}),
&json!({}),
unreachable_model_count,
);
assert_eq!(cost, 2);
}
#[test]
fn memory_remember_no_explicit_model_reads_registered_count() {
let cost = cost_unit_for_dispatch(
"memory.remember",
&json!({"content": "x"}),
&json!({}),
|| 4,
);
assert_eq!(cost, 5);
}
#[test]
fn update_and_recall_and_search_never_touch_registry() {
for verb in [
"update",
"memory.recall",
"knowledge.search",
"knowledge.compose",
] {
let cost =
cost_unit_for_dispatch(verb, &json!({}), &json!({}), unreachable_model_count);
assert_eq!(cost, 2, "verb {verb}: base_weight(1) + 1*1*1");
}
}
#[test]
fn compute_clamps_multiplication_overflow_to_i64_max() {
assert_eq!(compute(1, i64::MAX, 2, 1), i64::MAX);
}
#[test]
fn compute_clamps_addition_overflow_to_i64_max() {
assert_eq!(compute(i64::MAX, 1, 1, 1), i64::MAX);
}
#[test]
fn knowledge_index_extreme_total_clamps_to_i64_max() {
let cost = cost_unit_for_dispatch(
"knowledge.index",
&json!({}),
&json!({"total": i64::MAX}),
|| 2,
);
assert_eq!(cost, i64::MAX);
}
#[test]
fn resource_payload_shape_is_work_class_and_cost_unit_only() {
let payload = resource_payload(
"stats",
&json!({}),
&json!({}),
unreachable_model_count,
None,
);
assert_eq!(
payload,
json!({"work_class": "interactive", "cost_unit": 1}),
"resource payload must be exactly {{work_class, cost_unit}}, no request_id key \
when the caller supplied none"
);
}
#[test]
fn resource_payload_stamps_request_id_when_supplied() {
let payload = resource_payload(
"stats",
&json!({}),
&json!({}),
unreachable_model_count,
Some(42),
);
assert_eq!(
payload,
json!({"work_class": "interactive", "cost_unit": 1, "request_id": 42}),
);
}
#[test]
fn base_resource_payload_omits_request_id_when_absent() {
assert_eq!(
base_resource_payload(None),
json!({"work_class": "interactive"}),
);
}
#[test]
fn base_resource_payload_stamps_request_id_when_supplied() {
assert_eq!(
base_resource_payload(Some(7)),
json!({"work_class": "interactive", "request_id": 7}),
);
}
}