Skip to main content

lash_remote_protocol/
usage_activity.rs

1//! Token usage accounting and the streaming turn-activity event vocabulary.
2
3use schemars::JsonSchema;
4use serde::{Deserialize, Serialize};
5
6use crate::ensure_protocol_version;
7use crate::registry_errors::{RemoteProtocolError, require_non_empty};
8
9// Wire mirror of the runtime usage counters. This is a deliberately versioned
10// protocol boundary, kept independent of the internal types so the wire format
11// stays stable across internal refactors. The `From` converters in
12// `core_conversions::llm` destructure their `lash_core` source exhaustively
13// (no `..`), so adding a counter upstream is a compile error until it is mirrored
14// here too.
15#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
16pub struct RemoteUsage {
17    pub input_tokens: i64,
18    pub output_tokens: i64,
19    pub cache_read_input_tokens: i64,
20    pub cache_write_input_tokens: i64,
21    pub reasoning_output_tokens: i64,
22}
23
24impl RemoteUsage {
25    pub fn add(&mut self, other: &Self) {
26        self.input_tokens += other.input_tokens;
27        self.output_tokens += other.output_tokens;
28        self.cache_read_input_tokens += other.cache_read_input_tokens;
29        self.cache_write_input_tokens += other.cache_write_input_tokens;
30        self.reasoning_output_tokens += other.reasoning_output_tokens;
31    }
32}
33
34#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
35pub struct RemoteTokenLedgerEntry {
36    pub source: String,
37    pub model: String,
38    pub usage: RemoteUsage,
39}
40
41#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
42pub struct RemoteTurnActivity {
43    pub protocol_version: u32,
44    pub sequence: u64,
45    pub id: String,
46    pub correlation_id: String,
47    #[serde(flatten)]
48    pub event: RemoteTurnEvent,
49}
50
51impl RemoteTurnActivity {
52    pub fn validate(&self) -> Result<(), RemoteProtocolError> {
53        ensure_protocol_version(self.protocol_version)?;
54        require_non_empty("RemoteTurnActivity", "id", &self.id)?;
55        require_non_empty("RemoteTurnActivity", "correlation_id", &self.correlation_id)?;
56        if let RemoteTurnEvent::TurnInputApplied { applications } = &self.event {
57            for application in applications {
58                application.validate()?;
59            }
60        }
61        Ok(())
62    }
63}
64
65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
66#[serde(tag = "type", rename_all = "snake_case")]
67pub enum RemoteTurnEvent {
68    ModelRequestStarted {
69        protocol_iteration: usize,
70    },
71    AssistantProseDelta {
72        text: String,
73    },
74    ReasoningDelta {
75        text: String,
76    },
77    ModelAttemptReset {
78        assistant_prose_correlation_ids: Vec<String>,
79        reasoning_correlation_ids: Vec<String>,
80    },
81    CodeBlockStarted {
82        language: String,
83        code: String,
84        #[serde(default, skip_serializing_if = "Option::is_none")]
85        graph_key: Option<String>,
86    },
87    CodeBlockCompleted {
88        language: String,
89        output: String,
90        #[serde(default, skip_serializing_if = "Option::is_none")]
91        error: Option<String>,
92        success: bool,
93        duration_ms: u64,
94        tool_call_ids: Vec<String>,
95        #[serde(default, skip_serializing_if = "Option::is_none")]
96        graph_key: Option<String>,
97    },
98    ToolCallStarted {
99        #[serde(default, skip_serializing_if = "Option::is_none")]
100        call_id: Option<String>,
101        name: String,
102        args: serde_json::Value,
103        #[serde(default, skip_serializing_if = "Option::is_none")]
104        graph_key: Option<String>,
105        #[serde(default, skip_serializing_if = "Option::is_none")]
106        parent_call_id: Option<String>,
107    },
108    ToolCallCompleted {
109        #[serde(default, skip_serializing_if = "Option::is_none")]
110        call_id: Option<String>,
111        name: String,
112        args: serde_json::Value,
113        output: serde_json::Value,
114        duration_ms: u64,
115        #[serde(default, skip_serializing_if = "Option::is_none")]
116        graph_key: Option<String>,
117        #[serde(default, skip_serializing_if = "Option::is_none")]
118        parent_call_id: Option<String>,
119    },
120    FinalValue {
121        value: serde_json::Value,
122    },
123    ToolValue {
124        tool_name: String,
125        value: serde_json::Value,
126    },
127    Usage {
128        protocol_iteration: usize,
129        usage: RemoteUsage,
130        cumulative: RemoteUsage,
131    },
132    ChildUsage {
133        session_id: String,
134        source: String,
135        model: String,
136        protocol_iteration: usize,
137        usage: RemoteUsage,
138        cumulative: RemoteUsage,
139    },
140    RetryStatus {
141        wait_seconds: u64,
142        attempt: usize,
143        max_attempts: usize,
144        reason: String,
145    },
146    TurnInputApplied {
147        applications: Vec<crate::observations::RemoteTurnInputApplication>,
148    },
149    RuntimeDiagnostic {
150        kind: String,
151        data: serde_json::Value,
152    },
153    Error {
154        message: String,
155    },
156}