mcp_usage_core/
charge.rs

1//! The billing decision: what a request/response pair is worth.
2//!
3//! # The rule
4//!
5//! **Charge on terminal delivery, never on request.**
6//!
7//! Every consequence a meter needs falls out of that one sentence, with no chain
8//! correlation, no session, and no state:
9//!
10//! | Shape | Why it lands right |
11//! |---|---|
12//! | An MRTR chain of any length | Only its final round trip is `resultType: "complete"`, so it bills once |
13//! | An abandoned chain | Never reaches `complete`, so it bills nothing |
14//! | A task polled 300 times | Only the poll reporting `status: "completed"` delivered anything |
15//! | A task that fails or is cancelled | Terminal, but delivered nothing, so nothing is owed |
16//! | Discovery traffic | Never billable; it is the cost of connecting |
17//! | A protocol error | The server failed; charging for that is indefensible |
18//!
19//! This replaced an earlier design that tried to collapse MRTR chains by
20//! correlating their round trips. That design cannot be built by an intermediary:
21//! `requestState` is "an opaque string meaningful only to the server", which
22//! servers MUST integrity-protect with HMAC or AEAD, so the only correlator on
23//! the wire is encrypted under a key the meter does not hold. Charging on
24//! delivery needs no correlator, which is why it is both correct and cheaper.
25//!
26//! # Continuation state
27//!
28//! It does not consider whether the *request* was a continuation. Under the rule
29//! above that fact is irrelevant to billing, and requiring it would force a body
30//! peek on every `tools/call`. [`crate::peek::request`] still exposes it for the
31//! cache, because MRTR-derived results MUST NOT be cached.
32//!
33//! # Replay
34//!
35//! Two double-billing vectors exist, and they are treated differently on purpose.
36//!
37//! A **terminal task can be polled forever.** Terminal state "does not change", so
38//! every later `tasks/get` reports `completed` again. That is one unit of work
39//! reported many times, so [`Billable::idempotency_key`] carries the `taskId` and
40//! the caller MUST suppress repeats.
41//!
42//! A **replayed MRTR continuation** may produce a second `complete`. The spec
43//! bounds the replay window but explicitly does not guarantee single use. We bill
44//! it because the origin executed the tool a second time. Two completed executions
45//! produce two charges.
46
47use crate::peek::{ResponsePeek, ResultType};
48use crate::price::PriceBook;
49use crate::{Method, TaskStatus};
50
51/// The call being priced, as known from its (trusted) headers.
52///
53/// `name` must already be decoded from the `Mcp-Name` sentinel form.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct Call {
56    /// MCP method whose delivered work is being priced.
57    pub method: Method,
58    /// Decoded `Mcp-Name`, when the method carries one.
59    pub name: Option<String>,
60}
61
62impl Call {
63    /// Construct a classified call from trusted, decoded headers.
64    #[must_use]
65    pub fn new(method: Method, name: Option<String>) -> Self {
66        Self { method, name }
67    }
68}
69
70/// Non-identifying method category retained for a durable task.
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub enum TaskOriginKind {
73    /// A task created by `tools/call`.
74    ToolsCall,
75    /// A task created by `resources/read`.
76    ResourcesRead,
77    /// A task created by `prompts/get`.
78    PromptsGet,
79    /// A task created by an extension or unexpected method.
80    Other,
81}
82
83impl TaskOriginKind {
84    /// Reduce a method to a fixed category that cannot retain extension text.
85    #[must_use]
86    pub const fn from_method(method: &Method) -> Self {
87        match method {
88            Method::ToolsCall => Self::ToolsCall,
89            Method::ResourcesRead => Self::ResourcesRead,
90            Method::PromptsGet => Self::PromptsGet,
91            _ => Self::Other,
92        }
93    }
94
95    fn method(self) -> Method {
96        match self {
97            Self::ToolsCall => Method::ToolsCall,
98            Self::ResourcesRead => Method::ResourcesRead,
99            Self::PromptsGet => Method::PromptsGet,
100            Self::Other => Method::Other("[redacted-task-origin]".to_owned()),
101        }
102    }
103}
104
105/// Pre-priced durable-task attribution without the original name or URI.
106///
107/// Construct this when a task is created, then persist it under the task ID.
108/// The resolved units preserve named pricing while avoiding storage of a tool
109/// name, prompt name, resource URI, or extension method string.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub struct TaskAttribution {
112    origin_kind: TaskOriginKind,
113    units: u64,
114}
115
116impl TaskAttribution {
117    /// Construct an attribution from an already resolved method category and price.
118    #[must_use]
119    pub const fn new(origin_kind: TaskOriginKind, units: u64) -> Self {
120        Self { origin_kind, units }
121    }
122
123    /// Resolve a call against the current price book and discard identifying text.
124    #[must_use]
125    pub fn from_call(call: &Call, prices: &PriceBook) -> Self {
126        Self {
127            origin_kind: TaskOriginKind::from_method(&call.method),
128            units: prices.units_for(&call.method, call.name.as_deref()),
129        }
130    }
131
132    /// Fixed, non-identifying category of the call that created the task.
133    #[must_use]
134    pub const fn origin_kind(self) -> TaskOriginKind {
135        self.origin_kind
136    }
137
138    /// Units resolved when the task was created.
139    #[must_use]
140    pub const fn units(self) -> u64 {
141        self.units
142    }
143}
144
145/// A charge to record against a tenant.
146#[derive(Debug, Clone, PartialEq, Eq)]
147pub struct Billable {
148    /// Method that originally commissioned the delivered work.
149    pub method: Method,
150    /// Decoded tool, prompt, or resource name when available.
151    pub name: Option<String>,
152    /// Integer units to record with the billing backend.
153    pub units: u64,
154    /// When present, the caller MUST NOT record this charge more than once for
155    /// this key. Carried by task completions, whose terminal state can be
156    /// re-observed by any number of later polls.
157    pub idempotency_key: Option<String>,
158}
159
160/// Why a call was not billed. Kept specific so an operator answering "why is my
161/// bill lower than my request count" gets a real answer from the metrics.
162#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
163pub enum FreeReason {
164    /// `tools/list` and friends: the cost of connecting, not of work.
165    Discovery,
166    /// An interim `input_required` result. Nothing delivered yet.
167    InterimInputRequired,
168    /// A `CreateTaskResult`: work accepted, not performed.
169    TaskCreated,
170    /// A poll of a task still `working` or awaiting input.
171    TaskInProgress,
172    /// A task that reached a terminal state without delivering: failed, cancelled.
173    TaskNotDelivered,
174    /// `tasks/update` or `tasks/cancel`: driving a task, not receiving its output.
175    TaskDrive,
176    /// A completed task could not be joined to the call that created it.
177    MissingTaskAttribution,
178    /// A completed task omitted its required durable identifier, so repeat
179    /// polls could not be made idempotent.
180    MissingTaskId,
181    /// A `subscriptions/listen` stream. Not a discrete call; see the findings doc.
182    Subscription,
183    /// A JSON-RPC error. The server did not deliver.
184    ProtocolError,
185    /// A result shape this build does not recognize. Fails toward not charging.
186    UnrecognizedResult,
187}
188
189impl FreeReason {
190    /// Every reason, in a stable order.
191    ///
192    /// Callers index metric slots by position, so appending is safe and
193    /// reordering is not.
194    pub const ALL: [Self; 11] = [
195        Self::Discovery,
196        Self::InterimInputRequired,
197        Self::TaskCreated,
198        Self::TaskInProgress,
199        Self::TaskNotDelivered,
200        Self::TaskDrive,
201        Self::MissingTaskAttribution,
202        Self::MissingTaskId,
203        Self::Subscription,
204        Self::ProtocolError,
205        Self::UnrecognizedResult,
206    ];
207
208    /// The stable wire name.
209    ///
210    /// One spelling per variant, defined here so metrics, bindings and the
211    /// conformance vectors cannot drift into three different answers for the
212    /// same reason. Written out rather than derived from the identifier: a
213    /// rename must not silently change what an operator's dashboard or a
214    /// caller's branch is keyed on.
215    #[must_use]
216    pub const fn as_str(self) -> &'static str {
217        match self {
218            Self::Discovery => "discovery",
219            Self::InterimInputRequired => "interim_input_required",
220            Self::TaskCreated => "task_created",
221            Self::TaskInProgress => "task_in_progress",
222            Self::TaskNotDelivered => "task_not_delivered",
223            Self::TaskDrive => "task_drive",
224            Self::MissingTaskAttribution => "missing_task_attribution",
225            Self::MissingTaskId => "missing_task_id",
226            Self::Subscription => "subscription",
227            Self::ProtocolError => "protocol_error",
228            Self::UnrecognizedResult => "unrecognized_result",
229        }
230    }
231
232    /// Position in [`Self::ALL`], for indexing a fixed array of counters.
233    #[must_use]
234    pub const fn index(self) -> usize {
235        match self {
236            Self::Discovery => 0,
237            Self::InterimInputRequired => 1,
238            Self::TaskCreated => 2,
239            Self::TaskInProgress => 3,
240            Self::TaskNotDelivered => 4,
241            Self::TaskDrive => 5,
242            Self::MissingTaskAttribution => 6,
243            Self::MissingTaskId => 7,
244            Self::Subscription => 8,
245            Self::ProtocolError => 9,
246            Self::UnrecognizedResult => 10,
247        }
248    }
249}
250
251/// The verdict.
252#[derive(Debug, Clone, PartialEq, Eq)]
253pub enum Charge {
254    /// Delivered work that should be recorded.
255    Billable(Billable),
256    /// An exchange for which no usage is owed.
257    Free(FreeReason),
258}
259
260impl Charge {
261    /// Units owed, or zero when free.
262    #[must_use]
263    pub fn units(&self) -> u64 {
264        match self {
265            Self::Billable(b) => b.units,
266            Self::Free(_) => 0,
267        }
268    }
269
270    /// Whether the verdict records delivered work (including zero-priced work).
271    #[must_use]
272    pub const fn is_billable(&self) -> bool {
273        matches!(self, Self::Billable(_))
274    }
275
276    /// The idempotency key, when this charge carries one.
277    #[must_use]
278    pub fn idempotency_key(&self) -> Option<&str> {
279        match self {
280            Self::Billable(b) => b.idempotency_key.as_deref(),
281            Self::Free(_) => None,
282        }
283    }
284}
285
286/// Decide what a completed request/response exchange is worth.
287///
288/// Ordering matters: errors are checked before anything else so a failure is
289/// never billed regardless of which method produced it, and task handling
290/// precedes the generic `resultType` arm because a task poll legitimately carries
291/// `resultType: "complete"` while delivering nothing but a progress report.
292#[must_use]
293pub fn decide(call: &Call, response: &ResponsePeek, prices: &PriceBook) -> Charge {
294    decide_with_task_origin(call, response, prices, None)
295}
296
297enum TaskPricing<'a> {
298    Call(&'a Call),
299    Attribution(&'a TaskAttribution),
300}
301
302/// Decide an exchange while supplying the call that created a polled task.
303///
304/// A `tasks/get` request carries no `Mcp-Name`, so its terminal response cannot
305/// be priced correctly without the attribution captured from the original
306/// `tools/call`. Callers must persist that association when they observe a
307/// `resultType: "task"` response and pass it here on subsequent polls.
308#[must_use]
309pub fn decide_with_task_origin(
310    call: &Call,
311    response: &ResponsePeek,
312    prices: &PriceBook,
313    task_origin: Option<&Call>,
314) -> Charge {
315    decide_with_task_pricing(call, response, prices, task_origin.map(TaskPricing::Call))
316}
317
318/// Decide an exchange using pre-priced, non-identifying task attribution.
319///
320/// This is the privacy-preserving form for durable storage. The task's units
321/// are resolved when it is created, so a later terminal poll does not need the
322/// original name, URI, or extension method string.
323#[must_use]
324pub fn decide_with_task_attribution(
325    call: &Call,
326    response: &ResponsePeek,
327    prices: &PriceBook,
328    task_attribution: Option<&TaskAttribution>,
329) -> Charge {
330    decide_with_task_pricing(
331        call,
332        response,
333        prices,
334        task_attribution.map(TaskPricing::Attribution),
335    )
336}
337
338fn decide_with_task_pricing(
339    call: &Call,
340    response: &ResponsePeek,
341    prices: &PriceBook,
342    task_pricing: Option<TaskPricing<'_>>,
343) -> Charge {
344    if response.is_error {
345        return Charge::Free(FreeReason::ProtocolError);
346    }
347
348    if call.method.is_discovery() {
349        return Charge::Free(FreeReason::Discovery);
350    }
351
352    if matches!(call.method, Method::SubscriptionsListen) {
353        return Charge::Free(FreeReason::Subscription);
354    }
355
356    // A task handle accepts work but does not deliver it, regardless of the
357    // seed status fields carried beside the discriminant.
358    if matches!(response.result_type, ResultType::Task) {
359        return Charge::Free(FreeReason::TaskCreated);
360    }
361
362    // Task lifecycle. A poll's own resultType says the poll succeeded, which says
363    // nothing about the job, so the task status decides. The original call, not
364    // `tasks/get`, owns the price and attribution.
365    if matches!(call.method, Method::TasksGet)
366        && let Some(task) = &response.task
367    {
368        return match &task.status {
369            s if s.delivered() => {
370                let Some(task_id) = task.task_id.clone() else {
371                    return Charge::Free(FreeReason::MissingTaskId);
372                };
373                let Some(origin) = task_pricing else {
374                    return Charge::Free(FreeReason::MissingTaskAttribution);
375                };
376                let (method, name, units) = match origin {
377                    TaskPricing::Call(origin) => (
378                        origin.method.clone(),
379                        origin.name.clone(),
380                        prices.units_for(&origin.method, origin.name.as_deref()),
381                    ),
382                    TaskPricing::Attribution(attribution) => {
383                        (attribution.origin_kind.method(), None, attribution.units)
384                    }
385                };
386                Charge::Billable(Billable {
387                    method,
388                    name,
389                    units,
390                    // A terminal task reports `completed` to every later poll.
391                    idempotency_key: Some(task_id),
392                })
393            }
394            s if s.is_terminal() => Charge::Free(FreeReason::TaskNotDelivered),
395            TaskStatus::Working | TaskStatus::InputRequired => {
396                Charge::Free(FreeReason::TaskInProgress)
397            }
398            _ => Charge::Free(FreeReason::TaskInProgress),
399        };
400    }
401
402    if call.method.is_task_drive() {
403        return Charge::Free(FreeReason::TaskDrive);
404    }
405
406    match &response.result_type {
407        ResultType::Complete => Charge::Billable(Billable {
408            method: call.method.clone(),
409            name: call.name.clone(),
410            units: prices.units_for(&call.method, call.name.as_deref()),
411            idempotency_key: None,
412        }),
413        ResultType::InputRequired => Charge::Free(FreeReason::InterimInputRequired),
414        ResultType::Task => Charge::Free(FreeReason::TaskCreated),
415        ResultType::Unknown(_) | ResultType::Absent => Charge::Free(FreeReason::UnrecognizedResult),
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    #[test]
422    fn every_reason_has_a_distinct_name_and_a_matching_index() {
423        use super::FreeReason;
424        let mut names: Vec<&str> = FreeReason::ALL.iter().map(|r| r.as_str()).collect();
425        let total = names.len();
426        names.sort_unstable();
427        names.dedup();
428        assert_eq!(names.len(), total, "two reasons share a wire name");
429
430        for (position, reason) in FreeReason::ALL.iter().enumerate() {
431            assert_eq!(
432                reason.index(),
433                position,
434                "{} indexes a slot that is not its own",
435                reason.as_str()
436            );
437        }
438    }
439
440    use super::*;
441    use crate::peek::{TaskPeek, response as peek_response};
442    use serde_json::json;
443
444    fn call(method: Method, name: &str) -> Call {
445        Call::new(method, Some(name.to_owned()))
446    }
447
448    fn book() -> PriceBook {
449        PriceBook::flat(1)
450    }
451
452    fn complete() -> ResponsePeek {
453        peek_response(&json!({ "result": { "resultType": "complete", "content": [] } }))
454    }
455
456    fn input_required() -> ResponsePeek {
457        peek_response(&json!({
458            "result": {
459                "resultType": "input_required",
460                "inputRequests": { "k": { "method": "elicitation/create" } },
461                "requestState": "AEAD-protected blob"
462            }
463        }))
464    }
465
466    fn task_poll(status: &str, id: &str) -> ResponsePeek {
467        peek_response(&json!({
468            "result": { "resultType": "complete", "taskId": id, "status": status }
469        }))
470    }
471
472    #[test]
473    fn a_delivered_tool_call_bills_once() {
474        let charge = decide(
475            &call(Method::ToolsCall, "get_weather"),
476            &complete(),
477            &book(),
478        );
479        assert!(charge.is_billable());
480        assert_eq!(charge.units(), 1);
481        assert_eq!(charge.idempotency_key(), None);
482    }
483
484    #[test]
485    fn an_mrtr_chain_of_any_length_bills_exactly_once() {
486        // The property that motivated the whole crate. Rounds 1..n-1 are interim;
487        // only the last one delivers.
488        for rounds in 1..=10 {
489            let c = call(Method::ToolsCall, "get_weather");
490            let mut total = 0;
491            for round in 1..=rounds {
492                let response = if round == rounds {
493                    complete()
494                } else {
495                    input_required()
496                };
497                total += decide(&c, &response, &book()).units();
498            }
499            assert_eq!(total, 1, "a {rounds}-round chain billed {total} units");
500        }
501    }
502
503    #[test]
504    fn an_abandoned_chain_bills_nothing() {
505        let c = call(Method::ToolsCall, "get_weather");
506        let total: u64 = (0..5)
507            .map(|_| decide(&c, &input_required(), &book()).units())
508            .sum();
509        assert_eq!(total, 0);
510    }
511
512    #[test]
513    fn mrtr_applies_to_resources_read_and_prompts_get_too() {
514        for method in [Method::ResourcesRead, Method::PromptsGet] {
515            let c = call(method.clone(), "x");
516            assert_eq!(decide(&c, &input_required(), &book()).units(), 0);
517            assert_eq!(decide(&c, &complete(), &book()).units(), 1);
518        }
519    }
520
521    #[test]
522    fn a_task_polled_many_times_bills_once() {
523        let poll = Call::new(Method::TasksGet, None);
524        let origin = call(Method::ToolsCall, "long_job");
525        let mut total = 0;
526        let mut billed_keys: Vec<String> = Vec::new();
527
528        // 300 progress polls, then the completion.
529        for _ in 0..300 {
530            total += decide_with_task_origin(
531                &poll,
532                &task_poll("working", "tsk_1"),
533                &book(),
534                Some(&origin),
535            )
536            .units();
537        }
538        let final_charge = decide_with_task_origin(
539            &poll,
540            &task_poll("completed", "tsk_1"),
541            &book(),
542            Some(&origin),
543        );
544        if let Some(key) = final_charge.idempotency_key() {
545            billed_keys.push(key.to_owned());
546        }
547        total += final_charge.units();
548
549        assert_eq!(total, 1, "300 polls plus a completion billed {total} units");
550        assert_eq!(billed_keys, vec!["tsk_1".to_owned()]);
551    }
552
553    #[test]
554    fn a_completed_task_carries_its_id_so_repeat_polls_can_be_suppressed() {
555        // Terminal state does not change, so nothing stops a client polling a
556        // finished task forever. Core flags it; the caller enforces once-only.
557        let poll = Call::new(Method::TasksGet, None);
558        let origin = call(Method::ToolsCall, "long_job");
559        let first = decide_with_task_origin(
560            &poll,
561            &task_poll("completed", "tsk_9"),
562            &book(),
563            Some(&origin),
564        );
565        let second = decide_with_task_origin(
566            &poll,
567            &task_poll("completed", "tsk_9"),
568            &book(),
569            Some(&origin),
570        );
571
572        assert!(first.is_billable() && second.is_billable());
573        assert_eq!(first.idempotency_key(), Some("tsk_9"));
574        assert_eq!(first.idempotency_key(), second.idempotency_key());
575    }
576
577    #[test]
578    fn failed_and_cancelled_tasks_are_free() {
579        let c = Call::new(Method::TasksGet, None);
580        assert_eq!(
581            decide(&c, &task_poll("failed", "t"), &book()),
582            Charge::Free(FreeReason::TaskNotDelivered)
583        );
584        assert_eq!(
585            decide(&c, &task_poll("cancelled", "t"), &book()),
586            Charge::Free(FreeReason::TaskNotDelivered)
587        );
588    }
589
590    #[test]
591    fn a_task_awaiting_input_is_free() {
592        let c = Call::new(Method::TasksGet, None);
593        assert_eq!(
594            decide(&c, &task_poll("input_required", "t"), &book()),
595            Charge::Free(FreeReason::TaskInProgress)
596        );
597    }
598
599    #[test]
600    fn creating_a_task_is_not_delivering_it() {
601        let created = peek_response(&json!({
602            "result": { "resultType": "task", "taskId": "tsk_1", "status": "working" }
603        }));
604        assert_eq!(
605            decide(&call(Method::ToolsCall, "t"), &created, &book()),
606            Charge::Free(FreeReason::TaskCreated)
607        );
608    }
609
610    #[test]
611    fn task_drive_calls_are_free() {
612        let ack = peek_response(&json!({ "result": { "resultType": "complete" } }));
613        for method in [Method::TasksUpdate, Method::TasksCancel] {
614            assert_eq!(
615                decide(&Call::new(method, None), &ack, &book()),
616                Charge::Free(FreeReason::TaskDrive)
617            );
618        }
619    }
620
621    #[test]
622    fn discovery_is_never_billed() {
623        for method in [
624            Method::ToolsList,
625            Method::PromptsList,
626            Method::ResourcesList,
627            Method::ResourcesTemplatesList,
628            Method::ServerDiscover,
629        ] {
630            assert_eq!(
631                decide(&Call::new(method, None), &complete(), &book()),
632                Charge::Free(FreeReason::Discovery)
633            );
634        }
635    }
636
637    #[test]
638    fn subscription_streams_are_free() {
639        assert_eq!(
640            decide(
641                &Call::new(Method::SubscriptionsListen, None),
642                &complete(),
643                &book()
644            ),
645            Charge::Free(FreeReason::Subscription)
646        );
647    }
648
649    #[test]
650    fn errors_are_never_billed_whatever_the_method() {
651        let err = peek_response(&json!({ "error": { "code": -32020, "message": "mismatch" } }));
652        for method in [Method::ToolsCall, Method::TasksGet, Method::ResourcesRead] {
653            assert_eq!(
654                decide(&call(method, "x"), &err, &book()),
655                Charge::Free(FreeReason::ProtocolError)
656            );
657        }
658    }
659
660    #[test]
661    fn unrecognized_results_fail_toward_not_charging() {
662        let weird = peek_response(&json!({ "result": { "resultType": "some_future_shape" } }));
663        assert_eq!(
664            decide(&call(Method::ToolsCall, "t"), &weird, &book()),
665            Charge::Free(FreeReason::UnrecognizedResult)
666        );
667
668        let empty = peek_response(&json!({ "jsonrpc": "2.0" }));
669        assert_eq!(
670            decide(&call(Method::ToolsCall, "t"), &empty, &book()),
671            Charge::Free(FreeReason::UnrecognizedResult)
672        );
673    }
674
675    #[test]
676    fn price_book_drives_the_units() {
677        let prices = PriceBook::flat(1).with_name("expensive", 250);
678        assert_eq!(
679            decide(&call(Method::ToolsCall, "expensive"), &complete(), &prices).units(),
680            250
681        );
682        assert_eq!(
683            decide(&call(Method::ToolsCall, "cheap"), &complete(), &prices).units(),
684            1
685        );
686    }
687
688    #[test]
689    fn a_zero_priced_tool_still_reports_as_billable_with_zero_units() {
690        // Distinguishing "free by price" from "free by protocol shape" keeps the
691        // metrics honest about why a call cost nothing.
692        let prices = PriceBook::flat(5).with_name("loss_leader", 0);
693        let charge = decide(
694            &call(Method::ToolsCall, "loss_leader"),
695            &complete(),
696            &prices,
697        );
698        assert!(charge.is_billable());
699        assert_eq!(charge.units(), 0);
700    }
701
702    #[test]
703    fn an_unknown_task_status_is_not_billed() {
704        let paused = peek_response(&json!({
705            "result": { "resultType": "complete", "taskId": "t", "status": "paused" }
706        }));
707        assert!(!decide(&Call::new(Method::TasksGet, None), &paused, &book()).is_billable());
708    }
709
710    #[test]
711    fn a_task_without_an_id_fails_toward_not_charging() {
712        let peek = ResponsePeek {
713            result_type: ResultType::Complete,
714            is_error: false,
715            task: Some(TaskPeek {
716                task_id: None,
717                status: TaskStatus::Completed,
718            }),
719        };
720        let charge = decide_with_task_origin(
721            &Call::new(Method::TasksGet, None),
722            &peek,
723            &book(),
724            Some(&call(Method::ToolsCall, "j")),
725        );
726        assert_eq!(charge, Charge::Free(FreeReason::MissingTaskId));
727    }
728
729    #[test]
730    fn a_completed_task_without_origin_fails_toward_not_charging() {
731        let charge = decide(
732            &Call::new(Method::TasksGet, None),
733            &task_poll("completed", "tsk_orphan"),
734            &book(),
735        );
736        assert_eq!(charge, Charge::Free(FreeReason::MissingTaskAttribution));
737    }
738
739    #[test]
740    fn a_completed_task_uses_the_originating_tool_price() {
741        let prices = PriceBook::flat(1).with_name("expensive_job", 250);
742        let charge = decide_with_task_origin(
743            &Call::new(Method::TasksGet, None),
744            &task_poll("completed", "tsk_expensive"),
745            &prices,
746            Some(&call(Method::ToolsCall, "expensive_job")),
747        );
748        assert_eq!(charge.units(), 250);
749        let Charge::Billable(billable) = charge else {
750            panic!("completed task should be billable");
751        };
752        assert_eq!(billable.method, Method::ToolsCall);
753        assert_eq!(billable.name.as_deref(), Some("expensive_job"));
754    }
755
756    #[test]
757    fn prepriced_task_attribution_discards_identifying_names() {
758        let prices = PriceBook::flat(1).with_name("file:///private/customer-record", 250);
759        let origin = Call::new(
760            Method::ResourcesRead,
761            Some("file:///private/customer-record".to_owned()),
762        );
763        let attribution = TaskAttribution::from_call(&origin, &prices);
764
765        assert_eq!(attribution.units(), 250);
766        assert_eq!(attribution.origin_kind(), TaskOriginKind::ResourcesRead);
767        assert!(!format!("{attribution:?}").contains("customer-record"));
768
769        let charge = decide_with_task_attribution(
770            &Call::new(Method::TasksGet, None),
771            &task_poll("completed", "tsk_private"),
772            &prices,
773            Some(&attribution),
774        );
775        let Charge::Billable(billable) = charge else {
776            panic!("completed task should be billable");
777        };
778        assert_eq!(billable.units, 250);
779        assert_eq!(billable.method, Method::ResourcesRead);
780        assert_eq!(billable.name, None);
781    }
782}
mcp_usage_core/charge.rs

mcp_usage_core/
charge.rs
