| 1 | use serde::{Deserialize, Serialize}; |
| 2 | |
| 3 | /// Per-worker usage telemetry carried on `task_completed` events (#2974). |
| 4 | /// |
| 5 | /// Tokens come from the worker ledger (`AgentRunUsage`); `tool_calls` is the |
| 6 | /// worker's model/tool step count (`SubAgentResult::steps_taken`) and |
| 7 | /// `result_ref` points at the durable child artifact (transcript handle) so |
| 8 | /// consumers can fetch full output by reference instead of inline text. |
| 9 | /// Field names mirror `AgentRunUsage` so #4039 can render Tokens/Tools |
| 10 | /// columns without a remapping layer. |
| 11 | #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] |
| 12 | pub(super) struct WorkflowTaskUsage { |
| 13 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 14 | pub(super) input_tokens: Option<u64>, |
| 15 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 16 | pub(super) output_tokens: Option<u64>, |
| 17 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 18 | pub(super) total_tokens: Option<u64>, |
| 19 | /// Priced USD subtotal carried from the worker's immutable route audits, |
| 20 | /// in microdollars. Absence is unknown, never a zero-cost claim. |
| 21 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 22 | pub(super) cost_microusd: Option<u64>, |
| 23 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 24 | pub(super) tool_calls: Option<u32>, |
| 25 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 26 | pub(super) duration_ms: Option<u64>, |
| 27 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 28 | pub(super) result_ref: Option<String>, |
| 29 | /// Provenance of the token counts. This producer currently emits only |
| 30 | /// `provider_reported`; absent means unknown and must never render as zero |
| 31 | /// (#4039). |
| 32 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 33 | pub(super) token_source: Option<WorkflowTokenSource>, |
| 34 | } |
| 35 | |
| 36 | #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] |
| 37 | #[serde(rename_all = "snake_case")] |
| 38 | pub(super) enum WorkflowTokenSource { |
| 39 | ProviderReported, |
| 40 | } |
| 41 | |
| 42 | /// Run-wide usage totals reconciled from per-task telemetry, carried on |
| 43 | /// `run_completed` events and the persisted run record (#2974). |
| 44 | #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] |
| 45 | pub(super) struct WorkflowRunUsage { |
| 46 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 47 | pub(super) input_tokens: Option<u64>, |
| 48 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 49 | pub(super) output_tokens: Option<u64>, |
| 50 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 51 | pub(super) total_tokens: Option<u64>, |
| 52 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 53 | pub(super) cost_microusd: Option<u64>, |
| 54 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 55 | pub(super) tool_calls: Option<u64>, |
| 56 | /// Number of completed tasks that contributed telemetry. |
| 57 | #[serde(default)] |
| 58 | pub(super) tasks_reported: u64, |
| 59 | } |
| 60 | |
| 61 | impl WorkflowRunUsage { |
| 62 | pub(super) fn from_task(usage: &WorkflowTaskUsage) -> Self { |
| 63 | Self { |
| 64 | input_tokens: usage.input_tokens, |
| 65 | output_tokens: usage.output_tokens, |
| 66 | total_tokens: usage.total_tokens, |
| 67 | cost_microusd: usage.cost_microusd, |
| 68 | tool_calls: usage.tool_calls.map(u64::from), |
| 69 | tasks_reported: 1, |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | pub(super) fn add_task(&mut self, usage: &WorkflowTaskUsage) { |
| 74 | self.input_tokens = sum_optional_usage(self.input_tokens, usage.input_tokens); |
| 75 | self.output_tokens = sum_optional_usage(self.output_tokens, usage.output_tokens); |
| 76 | self.total_tokens = sum_optional_usage(self.total_tokens, usage.total_tokens); |
| 77 | self.cost_microusd = sum_optional_usage(self.cost_microusd, usage.cost_microusd); |
| 78 | self.tool_calls = sum_optional_usage(self.tool_calls, usage.tool_calls.map(u64::from)); |
| 79 | self.tasks_reported = self.tasks_reported.saturating_add(1); |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | pub(super) fn sum_optional_usage(left: Option<u64>, right: Option<u64>) -> Option<u64> { |
| 84 | match (left, right) { |
| 85 | (Some(left), Some(right)) => Some(left.saturating_add(right)), |
| 86 | (Some(value), None) | (None, Some(value)) => Some(value), |
| 87 | (None, None) => None, |
| 88 | } |
| 89 | } |
| 90 |