use crate::peek::{ResponsePeek, ResultType};
use crate::price::PriceBook;
use crate::{Method, TaskStatus};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Call {
pub method: Method,
pub name: Option<String>,
}
impl Call {
#[must_use]
pub fn new(method: Method, name: Option<String>) -> Self {
Self { method, name }
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskOriginKind {
ToolsCall,
ResourcesRead,
PromptsGet,
Other,
}
impl TaskOriginKind {
#[must_use]
pub const fn from_method(method: &Method) -> Self {
match method {
Method::ToolsCall => Self::ToolsCall,
Method::ResourcesRead => Self::ResourcesRead,
Method::PromptsGet => Self::PromptsGet,
_ => Self::Other,
}
}
fn method(self) -> Method {
match self {
Self::ToolsCall => Method::ToolsCall,
Self::ResourcesRead => Method::ResourcesRead,
Self::PromptsGet => Method::PromptsGet,
Self::Other => Method::Other("[redacted-task-origin]".to_owned()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TaskAttribution {
origin_kind: TaskOriginKind,
units: u64,
}
impl TaskAttribution {
#[must_use]
pub const fn new(origin_kind: TaskOriginKind, units: u64) -> Self {
Self { origin_kind, units }
}
#[must_use]
pub fn from_call(call: &Call, prices: &PriceBook) -> Self {
Self {
origin_kind: TaskOriginKind::from_method(&call.method),
units: prices.units_for(&call.method, call.name.as_deref()),
}
}
#[must_use]
pub const fn origin_kind(self) -> TaskOriginKind {
self.origin_kind
}
#[must_use]
pub const fn units(self) -> u64 {
self.units
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Billable {
pub method: Method,
pub name: Option<String>,
pub units: u64,
pub idempotency_key: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FreeReason {
Discovery,
InterimInputRequired,
TaskCreated,
TaskInProgress,
TaskNotDelivered,
TaskDrive,
MissingTaskAttribution,
MissingTaskId,
Subscription,
ProtocolError,
UnrecognizedResult,
}
impl FreeReason {
pub const ALL: [Self; 11] = [
Self::Discovery,
Self::InterimInputRequired,
Self::TaskCreated,
Self::TaskInProgress,
Self::TaskNotDelivered,
Self::TaskDrive,
Self::MissingTaskAttribution,
Self::MissingTaskId,
Self::Subscription,
Self::ProtocolError,
Self::UnrecognizedResult,
];
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Discovery => "discovery",
Self::InterimInputRequired => "interim_input_required",
Self::TaskCreated => "task_created",
Self::TaskInProgress => "task_in_progress",
Self::TaskNotDelivered => "task_not_delivered",
Self::TaskDrive => "task_drive",
Self::MissingTaskAttribution => "missing_task_attribution",
Self::MissingTaskId => "missing_task_id",
Self::Subscription => "subscription",
Self::ProtocolError => "protocol_error",
Self::UnrecognizedResult => "unrecognized_result",
}
}
#[must_use]
pub const fn index(self) -> usize {
match self {
Self::Discovery => 0,
Self::InterimInputRequired => 1,
Self::TaskCreated => 2,
Self::TaskInProgress => 3,
Self::TaskNotDelivered => 4,
Self::TaskDrive => 5,
Self::MissingTaskAttribution => 6,
Self::MissingTaskId => 7,
Self::Subscription => 8,
Self::ProtocolError => 9,
Self::UnrecognizedResult => 10,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Charge {
Billable(Billable),
Free(FreeReason),
}
impl Charge {
#[must_use]
pub fn units(&self) -> u64 {
match self {
Self::Billable(b) => b.units,
Self::Free(_) => 0,
}
}
#[must_use]
pub const fn is_billable(&self) -> bool {
matches!(self, Self::Billable(_))
}
#[must_use]
pub fn idempotency_key(&self) -> Option<&str> {
match self {
Self::Billable(b) => b.idempotency_key.as_deref(),
Self::Free(_) => None,
}
}
}
#[must_use]
pub fn decide(call: &Call, response: &ResponsePeek, prices: &PriceBook) -> Charge {
decide_with_task_origin(call, response, prices, None)
}
enum TaskPricing<'a> {
Call(&'a Call),
Attribution(&'a TaskAttribution),
}
#[must_use]
pub fn decide_with_task_origin(
call: &Call,
response: &ResponsePeek,
prices: &PriceBook,
task_origin: Option<&Call>,
) -> Charge {
decide_with_task_pricing(call, response, prices, task_origin.map(TaskPricing::Call))
}
#[must_use]
pub fn decide_with_task_attribution(
call: &Call,
response: &ResponsePeek,
prices: &PriceBook,
task_attribution: Option<&TaskAttribution>,
) -> Charge {
decide_with_task_pricing(
call,
response,
prices,
task_attribution.map(TaskPricing::Attribution),
)
}
fn decide_with_task_pricing(
call: &Call,
response: &ResponsePeek,
prices: &PriceBook,
task_pricing: Option<TaskPricing<'_>>,
) -> Charge {
if response.is_error {
return Charge::Free(FreeReason::ProtocolError);
}
if call.method.is_discovery() {
return Charge::Free(FreeReason::Discovery);
}
if matches!(call.method, Method::SubscriptionsListen) {
return Charge::Free(FreeReason::Subscription);
}
if matches!(response.result_type, ResultType::Task) {
return Charge::Free(FreeReason::TaskCreated);
}
if matches!(call.method, Method::TasksGet)
&& let Some(task) = &response.task
{
return match &task.status {
s if s.delivered() => {
let Some(task_id) = task.task_id.clone() else {
return Charge::Free(FreeReason::MissingTaskId);
};
let Some(origin) = task_pricing else {
return Charge::Free(FreeReason::MissingTaskAttribution);
};
let (method, name, units) = match origin {
TaskPricing::Call(origin) => (
origin.method.clone(),
origin.name.clone(),
prices.units_for(&origin.method, origin.name.as_deref()),
),
TaskPricing::Attribution(attribution) => {
(attribution.origin_kind.method(), None, attribution.units)
}
};
Charge::Billable(Billable {
method,
name,
units,
idempotency_key: Some(task_id),
})
}
s if s.is_terminal() => Charge::Free(FreeReason::TaskNotDelivered),
TaskStatus::Working | TaskStatus::InputRequired => {
Charge::Free(FreeReason::TaskInProgress)
}
_ => Charge::Free(FreeReason::TaskInProgress),
};
}
if call.method.is_task_drive() {
return Charge::Free(FreeReason::TaskDrive);
}
match &response.result_type {
ResultType::Complete => Charge::Billable(Billable {
method: call.method.clone(),
name: call.name.clone(),
units: prices.units_for(&call.method, call.name.as_deref()),
idempotency_key: None,
}),
ResultType::InputRequired => Charge::Free(FreeReason::InterimInputRequired),
ResultType::Task => Charge::Free(FreeReason::TaskCreated),
ResultType::Unknown(_) | ResultType::Absent => Charge::Free(FreeReason::UnrecognizedResult),
}
}
#[cfg(test)]
mod tests {
#[test]
fn every_reason_has_a_distinct_name_and_a_matching_index() {
use super::FreeReason;
let mut names: Vec<&str> = FreeReason::ALL.iter().map(|r| r.as_str()).collect();
let total = names.len();
names.sort_unstable();
names.dedup();
assert_eq!(names.len(), total, "two reasons share a wire name");
for (position, reason) in FreeReason::ALL.iter().enumerate() {
assert_eq!(
reason.index(),
position,
"{} indexes a slot that is not its own",
reason.as_str()
);
}
}
use super::*;
use crate::peek::{TaskPeek, response as peek_response};
use serde_json::json;
fn call(method: Method, name: &str) -> Call {
Call::new(method, Some(name.to_owned()))
}
fn book() -> PriceBook {
PriceBook::flat(1)
}
fn complete() -> ResponsePeek {
peek_response(&json!({ "result": { "resultType": "complete", "content": [] } }))
}
fn input_required() -> ResponsePeek {
peek_response(&json!({
"result": {
"resultType": "input_required",
"inputRequests": { "k": { "method": "elicitation/create" } },
"requestState": "AEAD-protected blob"
}
}))
}
fn task_poll(status: &str, id: &str) -> ResponsePeek {
peek_response(&json!({
"result": { "resultType": "complete", "taskId": id, "status": status }
}))
}
#[test]
fn a_delivered_tool_call_bills_once() {
let charge = decide(
&call(Method::ToolsCall, "get_weather"),
&complete(),
&book(),
);
assert!(charge.is_billable());
assert_eq!(charge.units(), 1);
assert_eq!(charge.idempotency_key(), None);
}
#[test]
fn an_mrtr_chain_of_any_length_bills_exactly_once() {
for rounds in 1..=10 {
let c = call(Method::ToolsCall, "get_weather");
let mut total = 0;
for round in 1..=rounds {
let response = if round == rounds {
complete()
} else {
input_required()
};
total += decide(&c, &response, &book()).units();
}
assert_eq!(total, 1, "a {rounds}-round chain billed {total} units");
}
}
#[test]
fn an_abandoned_chain_bills_nothing() {
let c = call(Method::ToolsCall, "get_weather");
let total: u64 = (0..5)
.map(|_| decide(&c, &input_required(), &book()).units())
.sum();
assert_eq!(total, 0);
}
#[test]
fn mrtr_applies_to_resources_read_and_prompts_get_too() {
for method in [Method::ResourcesRead, Method::PromptsGet] {
let c = call(method.clone(), "x");
assert_eq!(decide(&c, &input_required(), &book()).units(), 0);
assert_eq!(decide(&c, &complete(), &book()).units(), 1);
}
}
#[test]
fn a_task_polled_many_times_bills_once() {
let poll = Call::new(Method::TasksGet, None);
let origin = call(Method::ToolsCall, "long_job");
let mut total = 0;
let mut billed_keys: Vec<String> = Vec::new();
for _ in 0..300 {
total += decide_with_task_origin(
&poll,
&task_poll("working", "tsk_1"),
&book(),
Some(&origin),
)
.units();
}
let final_charge = decide_with_task_origin(
&poll,
&task_poll("completed", "tsk_1"),
&book(),
Some(&origin),
);
if let Some(key) = final_charge.idempotency_key() {
billed_keys.push(key.to_owned());
}
total += final_charge.units();
assert_eq!(total, 1, "300 polls plus a completion billed {total} units");
assert_eq!(billed_keys, vec!["tsk_1".to_owned()]);
}
#[test]
fn a_completed_task_carries_its_id_so_repeat_polls_can_be_suppressed() {
let poll = Call::new(Method::TasksGet, None);
let origin = call(Method::ToolsCall, "long_job");
let first = decide_with_task_origin(
&poll,
&task_poll("completed", "tsk_9"),
&book(),
Some(&origin),
);
let second = decide_with_task_origin(
&poll,
&task_poll("completed", "tsk_9"),
&book(),
Some(&origin),
);
assert!(first.is_billable() && second.is_billable());
assert_eq!(first.idempotency_key(), Some("tsk_9"));
assert_eq!(first.idempotency_key(), second.idempotency_key());
}
#[test]
fn failed_and_cancelled_tasks_are_free() {
let c = Call::new(Method::TasksGet, None);
assert_eq!(
decide(&c, &task_poll("failed", "t"), &book()),
Charge::Free(FreeReason::TaskNotDelivered)
);
assert_eq!(
decide(&c, &task_poll("cancelled", "t"), &book()),
Charge::Free(FreeReason::TaskNotDelivered)
);
}
#[test]
fn a_task_awaiting_input_is_free() {
let c = Call::new(Method::TasksGet, None);
assert_eq!(
decide(&c, &task_poll("input_required", "t"), &book()),
Charge::Free(FreeReason::TaskInProgress)
);
}
#[test]
fn creating_a_task_is_not_delivering_it() {
let created = peek_response(&json!({
"result": { "resultType": "task", "taskId": "tsk_1", "status": "working" }
}));
assert_eq!(
decide(&call(Method::ToolsCall, "t"), &created, &book()),
Charge::Free(FreeReason::TaskCreated)
);
}
#[test]
fn task_drive_calls_are_free() {
let ack = peek_response(&json!({ "result": { "resultType": "complete" } }));
for method in [Method::TasksUpdate, Method::TasksCancel] {
assert_eq!(
decide(&Call::new(method, None), &ack, &book()),
Charge::Free(FreeReason::TaskDrive)
);
}
}
#[test]
fn discovery_is_never_billed() {
for method in [
Method::ToolsList,
Method::PromptsList,
Method::ResourcesList,
Method::ResourcesTemplatesList,
Method::ServerDiscover,
] {
assert_eq!(
decide(&Call::new(method, None), &complete(), &book()),
Charge::Free(FreeReason::Discovery)
);
}
}
#[test]
fn subscription_streams_are_free() {
assert_eq!(
decide(
&Call::new(Method::SubscriptionsListen, None),
&complete(),
&book()
),
Charge::Free(FreeReason::Subscription)
);
}
#[test]
fn errors_are_never_billed_whatever_the_method() {
let err = peek_response(&json!({ "error": { "code": -32020, "message": "mismatch" } }));
for method in [Method::ToolsCall, Method::TasksGet, Method::ResourcesRead] {
assert_eq!(
decide(&call(method, "x"), &err, &book()),
Charge::Free(FreeReason::ProtocolError)
);
}
}
#[test]
fn unrecognized_results_fail_toward_not_charging() {
let weird = peek_response(&json!({ "result": { "resultType": "some_future_shape" } }));
assert_eq!(
decide(&call(Method::ToolsCall, "t"), &weird, &book()),
Charge::Free(FreeReason::UnrecognizedResult)
);
let empty = peek_response(&json!({ "jsonrpc": "2.0" }));
assert_eq!(
decide(&call(Method::ToolsCall, "t"), &empty, &book()),
Charge::Free(FreeReason::UnrecognizedResult)
);
}
#[test]
fn price_book_drives_the_units() {
let prices = PriceBook::flat(1).with_name("expensive", 250);
assert_eq!(
decide(&call(Method::ToolsCall, "expensive"), &complete(), &prices).units(),
250
);
assert_eq!(
decide(&call(Method::ToolsCall, "cheap"), &complete(), &prices).units(),
1
);
}
#[test]
fn a_zero_priced_tool_still_reports_as_billable_with_zero_units() {
let prices = PriceBook::flat(5).with_name("loss_leader", 0);
let charge = decide(
&call(Method::ToolsCall, "loss_leader"),
&complete(),
&prices,
);
assert!(charge.is_billable());
assert_eq!(charge.units(), 0);
}
#[test]
fn an_unknown_task_status_is_not_billed() {
let paused = peek_response(&json!({
"result": { "resultType": "complete", "taskId": "t", "status": "paused" }
}));
assert!(!decide(&Call::new(Method::TasksGet, None), &paused, &book()).is_billable());
}
#[test]
fn a_task_without_an_id_fails_toward_not_charging() {
let peek = ResponsePeek {
result_type: ResultType::Complete,
is_error: false,
task: Some(TaskPeek {
task_id: None,
status: TaskStatus::Completed,
}),
};
let charge = decide_with_task_origin(
&Call::new(Method::TasksGet, None),
&peek,
&book(),
Some(&call(Method::ToolsCall, "j")),
);
assert_eq!(charge, Charge::Free(FreeReason::MissingTaskId));
}
#[test]
fn a_completed_task_without_origin_fails_toward_not_charging() {
let charge = decide(
&Call::new(Method::TasksGet, None),
&task_poll("completed", "tsk_orphan"),
&book(),
);
assert_eq!(charge, Charge::Free(FreeReason::MissingTaskAttribution));
}
#[test]
fn a_completed_task_uses_the_originating_tool_price() {
let prices = PriceBook::flat(1).with_name("expensive_job", 250);
let charge = decide_with_task_origin(
&Call::new(Method::TasksGet, None),
&task_poll("completed", "tsk_expensive"),
&prices,
Some(&call(Method::ToolsCall, "expensive_job")),
);
assert_eq!(charge.units(), 250);
let Charge::Billable(billable) = charge else {
panic!("completed task should be billable");
};
assert_eq!(billable.method, Method::ToolsCall);
assert_eq!(billable.name.as_deref(), Some("expensive_job"));
}
#[test]
fn prepriced_task_attribution_discards_identifying_names() {
let prices = PriceBook::flat(1).with_name("file:///private/customer-record", 250);
let origin = Call::new(
Method::ResourcesRead,
Some("file:///private/customer-record".to_owned()),
);
let attribution = TaskAttribution::from_call(&origin, &prices);
assert_eq!(attribution.units(), 250);
assert_eq!(attribution.origin_kind(), TaskOriginKind::ResourcesRead);
assert!(!format!("{attribution:?}").contains("customer-record"));
let charge = decide_with_task_attribution(
&Call::new(Method::TasksGet, None),
&task_poll("completed", "tsk_private"),
&prices,
Some(&attribution),
);
let Charge::Billable(billable) = charge else {
panic!("completed task should be billable");
};
assert_eq!(billable.units, 250);
assert_eq!(billable.method, Method::ResourcesRead);
assert_eq!(billable.name, None);
}
}