lean_ctx/core/
agent_attribution.rs

1//! ETPAO Attribution for multi-agent chains (P11 / DIM 4).
2//!
3//! Attribution rule: agent costs count only when the overall chain outcome is
4//! accepted. This prevents double-counting and ensures that failed/redundant
5//! branches don't inflate savings or costs.
6//!
7//! ETPAO = Effective Token Per Accepted Outcome
8
9use std::collections::BTreeMap;
10
11use serde::{Deserialize, Serialize};
12
13pub const ATTRIBUTION_SCHEMA_VERSION: u16 = 1;
14
15#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
16#[serde(rename_all = "snake_case")]
17pub enum OutcomeVerdict {
18    Pending,
19    Accepted,
20    Rejected,
21    Partial,
22}
23
24#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
25pub struct AgentCostRecord {
26    pub agent_id: String,
27    pub node_id: String,
28    pub tokens_consumed: u64,
29    pub cost_micros: u64,
30    pub is_on_accepted_path: bool,
31}
32
33#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
34pub struct ChainAttribution {
35    pub schema_version: u16,
36    pub chain_id: String,
37    pub outcome_verdict: OutcomeVerdict,
38    pub total_tokens_all_agents: u64,
39    pub total_cost_all_agents: u64,
40    pub effective_tokens_accepted: u64,
41    pub effective_cost_accepted: u64,
42    pub waste_tokens: u64,
43    pub waste_cost: u64,
44    pub agent_contributions: Vec<AgentContribution>,
45}
46
47#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
48pub struct AgentContribution {
49    pub agent_id: String,
50    pub tokens: u64,
51    pub cost_micros: u64,
52    pub fraction_of_accepted: f64,
53    pub on_accepted_path: bool,
54}
55
56/// Tracks costs per agent/node and computes ETPAO attribution once the chain
57/// outcome is known.
58pub struct AttributionTracker {
59    chain_id: String,
60    records: BTreeMap<String, AgentCostRecord>,
61    outcome: OutcomeVerdict,
62}
63
64impl AttributionTracker {
65    #[must_use]
66    pub fn new(chain_id: String) -> Self {
67        Self {
68            chain_id,
69            records: BTreeMap::new(),
70            outcome: OutcomeVerdict::Pending,
71        }
72    }
73
74    /// Record cost for an agent at a specific work node.
75    pub fn record_cost(
76        &mut self,
77        agent_id: String,
78        node_id: String,
79        tokens: u64,
80        cost_micros: u64,
81    ) {
82        self.records
83            .entry(node_id.clone())
84            .and_modify(|r| {
85                r.tokens_consumed = r.tokens_consumed.saturating_add(tokens);
86                r.cost_micros = r.cost_micros.saturating_add(cost_micros);
87            })
88            .or_insert(AgentCostRecord {
89                agent_id,
90                node_id,
91                tokens_consumed: tokens,
92                cost_micros,
93                is_on_accepted_path: false,
94            });
95    }
96
97    /// Mark which nodes are on the accepted outcome path.
98    pub fn mark_accepted_path(&mut self, node_ids: &[String]) {
99        for id in node_ids {
100            if let Some(record) = self.records.get_mut(id) {
101                record.is_on_accepted_path = true;
102            }
103        }
104    }
105
106    /// Set the final chain outcome verdict.
107    pub fn set_outcome(&mut self, verdict: OutcomeVerdict) {
108        self.outcome = verdict;
109    }
110
111    /// Compute the final ETPAO attribution. Costs on rejected/redundant paths
112    /// are classified as waste; only accepted-path costs count toward the
113    /// effective metric.
114    pub fn compute_attribution(&self) -> ChainAttribution {
115        let total_tokens: u64 = self.records.values().map(|r| r.tokens_consumed).sum();
116        let total_cost: u64 = self.records.values().map(|r| r.cost_micros).sum();
117
118        let (effective_tokens, effective_cost) = match self.outcome {
119            OutcomeVerdict::Accepted | OutcomeVerdict::Partial => {
120                let t: u64 = self
121                    .records
122                    .values()
123                    .filter(|r| r.is_on_accepted_path)
124                    .map(|r| r.tokens_consumed)
125                    .sum();
126                let c: u64 = self
127                    .records
128                    .values()
129                    .filter(|r| r.is_on_accepted_path)
130                    .map(|r| r.cost_micros)
131                    .sum();
132                (t, c)
133            }
134            OutcomeVerdict::Rejected | OutcomeVerdict::Pending => (0, 0),
135        };
136
137        let waste_tokens = total_tokens.saturating_sub(effective_tokens);
138        let waste_cost = total_cost.saturating_sub(effective_cost);
139
140        let mut contributions: BTreeMap<String, (u64, u64, bool)> = BTreeMap::new();
141        for record in self.records.values() {
142            let entry = contributions.entry(record.agent_id.clone()).or_default();
143            entry.0 = entry.0.saturating_add(record.tokens_consumed);
144            entry.1 = entry.1.saturating_add(record.cost_micros);
145            if record.is_on_accepted_path {
146                entry.2 = true;
147            }
148        }
149
150        let agent_contributions: Vec<AgentContribution> = contributions
151            .into_iter()
152            .map(|(agent_id, (tokens, cost, on_path))| {
153                let fraction = if effective_tokens > 0 && on_path {
154                    tokens as f64 / effective_tokens as f64
155                } else {
156                    0.0
157                };
158                AgentContribution {
159                    agent_id,
160                    tokens,
161                    cost_micros: cost,
162                    fraction_of_accepted: fraction,
163                    on_accepted_path: on_path,
164                }
165            })
166            .collect();
167
168        ChainAttribution {
169            schema_version: ATTRIBUTION_SCHEMA_VERSION,
170            chain_id: self.chain_id.clone(),
171            outcome_verdict: self.outcome,
172            total_tokens_all_agents: total_tokens,
173            total_cost_all_agents: total_cost,
174            effective_tokens_accepted: effective_tokens,
175            effective_cost_accepted: effective_cost,
176            waste_tokens,
177            waste_cost,
178            agent_contributions,
179        }
180    }
181
182    pub fn outcome(&self) -> OutcomeVerdict {
183        self.outcome
184    }
185}
186
187// ─── Tests ───────────────────────────────────────────────────────────────────
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn accepted_outcome_attributes_only_accepted_path() {
195        let mut tracker = AttributionTracker::new("chain:1".into());
196        tracker.record_cost("agent-a".into(), "node-root".into(), 500, 100);
197        tracker.record_cost("agent-b".into(), "node-child-1".into(), 300, 60);
198        tracker.record_cost("agent-c".into(), "node-child-2".into(), 200, 40);
199
200        tracker.mark_accepted_path(&["node-root".into(), "node-child-1".into()]);
201        tracker.set_outcome(OutcomeVerdict::Accepted);
202
203        let attr = tracker.compute_attribution();
204        assert_eq!(attr.effective_tokens_accepted, 800);
205        assert_eq!(attr.effective_cost_accepted, 160);
206        assert_eq!(attr.waste_tokens, 200);
207        assert_eq!(attr.waste_cost, 40);
208        assert_eq!(attr.total_tokens_all_agents, 1000);
209    }
210
211    #[test]
212    fn rejected_outcome_counts_all_as_waste() {
213        let mut tracker = AttributionTracker::new("chain:2".into());
214        tracker.record_cost("agent-a".into(), "n1".into(), 500, 100);
215        tracker.record_cost("agent-b".into(), "n2".into(), 300, 60);
216        tracker.mark_accepted_path(&["n1".into()]);
217        tracker.set_outcome(OutcomeVerdict::Rejected);
218
219        let attr = tracker.compute_attribution();
220        assert_eq!(attr.effective_tokens_accepted, 0);
221        assert_eq!(attr.waste_tokens, 800);
222    }
223
224    #[test]
225    fn pending_outcome_attributes_nothing() {
226        let mut tracker = AttributionTracker::new("chain:3".into());
227        tracker.record_cost("agent-a".into(), "n1".into(), 100, 20);
228        let attr = tracker.compute_attribution();
229        assert_eq!(attr.outcome_verdict, OutcomeVerdict::Pending);
230        assert_eq!(attr.effective_tokens_accepted, 0);
231    }
232
233    #[test]
234    fn agent_contribution_fractions_sum_to_one() {
235        let mut tracker = AttributionTracker::new("chain:4".into());
236        tracker.record_cost("agent-a".into(), "n1".into(), 600, 120);
237        tracker.record_cost("agent-b".into(), "n2".into(), 400, 80);
238        tracker.mark_accepted_path(&["n1".into(), "n2".into()]);
239        tracker.set_outcome(OutcomeVerdict::Accepted);
240
241        let attr = tracker.compute_attribution();
242        let sum: f64 = attr
243            .agent_contributions
244            .iter()
245            .filter(|c| c.on_accepted_path)
246            .map(|c| c.fraction_of_accepted)
247            .sum();
248        assert!((sum - 1.0).abs() < 1e-10);
249    }
250
251    #[test]
252    fn incremental_cost_recording() {
253        let mut tracker = AttributionTracker::new("chain:5".into());
254        tracker.record_cost("agent-a".into(), "n1".into(), 100, 20);
255        tracker.record_cost("agent-a".into(), "n1".into(), 50, 10);
256        tracker.mark_accepted_path(&["n1".into()]);
257        tracker.set_outcome(OutcomeVerdict::Accepted);
258
259        let attr = tracker.compute_attribution();
260        assert_eq!(attr.effective_tokens_accepted, 150);
261        assert_eq!(attr.effective_cost_accepted, 30);
262    }
263}
lean_ctx/core/agent_attribution.rs

lean_ctx/core/
agent_attribution.rs
