a3s_code_core/harness_evidence/
usage.rs

1use super::digest::{measure, require_digest, require_optional_digest};
2use super::{HarnessEvidenceError, ModelInputSnapshotV1};
3use crate::llm::TokenUsage;
4use serde::{Deserialize, Serialize};
5
6pub const MODEL_USAGE_SNAPSHOT_V1_SCHEMA: &str = "a3s.code.model-usage-snapshot.v1";
7const MODEL_USAGE_SNAPSHOT_DOMAIN: &str = "a3s.code.model-usage-snapshot.v1";
8
9#[derive(Clone, Debug)]
10pub(crate) struct ModelUsageBinding {
11    call_sequence: u64,
12    input_snapshot_digest: String,
13    estimated_prompt_tokens: usize,
14    tool_results: ToolResultContextUsageV1,
15}
16
17impl ModelUsageBinding {
18    pub(crate) fn from_input(
19        input: &ModelInputSnapshotV1,
20        tool_results: ToolResultContextUsageV1,
21    ) -> Self {
22        Self {
23            call_sequence: input.call_sequence,
24            input_snapshot_digest: input.snapshot_digest.clone(),
25            estimated_prompt_tokens: input.estimated_prompt_tokens,
26            tool_results,
27        }
28    }
29
30    pub(crate) fn call_sequence(&self) -> u64 {
31        self.call_sequence
32    }
33}
34
35/// Bounded measurements of Tool-result content visible to one model call.
36#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
37#[serde(rename_all = "camelCase", deny_unknown_fields)]
38pub struct ToolResultContextUsageV1 {
39    pub total_count: usize,
40    pub unique_count: usize,
41    pub repeated_count: usize,
42    pub content_bytes: u64,
43    pub repeated_content_bytes: u64,
44    pub estimated_tokens: usize,
45    pub repeated_estimated_tokens: usize,
46    pub contents_digest: Option<String>,
47    pub repeated_contents_digest: Option<String>,
48}
49
50impl ToolResultContextUsageV1 {
51    pub fn validate(&self) -> Result<(), HarnessEvidenceError> {
52        if self.unique_count > self.total_count
53            || self.repeated_count != self.total_count.saturating_sub(self.unique_count)
54            || (self.total_count > 0 && self.unique_count == 0)
55        {
56            return Err(HarnessEvidenceError::InvalidContents(
57                "unique and repeated Tool-result counts partition Tool results",
58            ));
59        }
60        let no_tool_results = self.total_count == 0;
61        if no_tool_results
62            != (self.unique_count == 0
63                && self.content_bytes == 0
64                && self.estimated_tokens == 0
65                && self.contents_digest.is_none())
66        {
67            return Err(HarnessEvidenceError::InvalidContents(
68                "Tool-result count, content usage, and digest agree",
69            ));
70        }
71        let no_repeated_tool_results = self.repeated_count == 0;
72        if no_repeated_tool_results
73            != (self.repeated_content_bytes == 0
74                && self.repeated_estimated_tokens == 0
75                && self.repeated_contents_digest.is_none())
76        {
77            return Err(HarnessEvidenceError::InvalidContents(
78                "repeated Tool-result count, content usage, and digest agree",
79            ));
80        }
81        if self.repeated_content_bytes > self.content_bytes
82            || self.repeated_estimated_tokens > self.estimated_tokens
83        {
84            return Err(HarnessEvidenceError::InvalidContents(
85                "repeated Tool-result usage is bounded by total Tool-result usage",
86            ));
87        }
88        require_optional_digest("contents_digest", self.contents_digest.as_deref())?;
89        require_optional_digest(
90            "repeated_contents_digest",
91            self.repeated_contents_digest.as_deref(),
92        )?;
93        Ok(())
94    }
95}
96
97/// Immutable per-call correlation between Code's prompt estimate and the
98/// normalized usage report returned by an [`crate::llm::LlmClient`].
99///
100/// These values are context diagnostics, not a billing ledger. A client that
101/// cannot observe provider usage may return zeroes, which are preserved rather
102/// than replaced with Code's estimate.
103#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
104#[serde(rename_all = "camelCase", deny_unknown_fields)]
105pub struct ModelUsageSnapshotV1 {
106    pub schema: String,
107    pub call_sequence: u64,
108    pub input_snapshot_digest: String,
109    pub estimated_prompt_tokens: usize,
110    pub reported_prompt_tokens: usize,
111    pub reported_completion_tokens: usize,
112    pub reported_total_tokens: usize,
113    pub reported_cache_read_tokens: Option<usize>,
114    pub reported_cache_write_tokens: Option<usize>,
115    pub tool_results: ToolResultContextUsageV1,
116    pub snapshot_digest: String,
117}
118
119impl ModelUsageSnapshotV1 {
120    pub fn from_input(
121        input: &ModelInputSnapshotV1,
122        tool_results: &ToolResultContextUsageV1,
123        usage: &TokenUsage,
124    ) -> Result<Self, HarnessEvidenceError> {
125        input.validate()?;
126        tool_results.validate()?;
127        if tool_results.total_count != input.tool_result_count {
128            return Err(HarnessEvidenceError::InvalidContents(
129                "usage and input Tool-result counts agree",
130            ));
131        }
132        Self::from_binding(
133            &ModelUsageBinding::from_input(input, tool_results.clone()),
134            usage,
135        )
136    }
137
138    pub(crate) fn from_binding(
139        binding: &ModelUsageBinding,
140        usage: &TokenUsage,
141    ) -> Result<Self, HarnessEvidenceError> {
142        binding.tool_results.validate()?;
143        let mut snapshot = Self {
144            schema: MODEL_USAGE_SNAPSHOT_V1_SCHEMA.to_string(),
145            call_sequence: binding.call_sequence,
146            input_snapshot_digest: binding.input_snapshot_digest.clone(),
147            estimated_prompt_tokens: binding.estimated_prompt_tokens,
148            reported_prompt_tokens: usage.prompt_tokens,
149            reported_completion_tokens: usage.completion_tokens,
150            reported_total_tokens: usage.total_tokens,
151            reported_cache_read_tokens: usage.cache_read_tokens,
152            reported_cache_write_tokens: usage.cache_write_tokens,
153            tool_results: binding.tool_results.clone(),
154            snapshot_digest: String::new(),
155        };
156        snapshot.snapshot_digest = snapshot.expected_digest()?;
157        Ok(snapshot)
158    }
159
160    pub fn validate(&self) -> Result<(), HarnessEvidenceError> {
161        if self.schema != MODEL_USAGE_SNAPSHOT_V1_SCHEMA {
162            return Err(HarnessEvidenceError::UnsupportedSchema);
163        }
164        if self.call_sequence == 0 {
165            return Err(HarnessEvidenceError::InvalidContents(
166                "call_sequence is positive",
167            ));
168        }
169        require_digest("input_snapshot_digest", &self.input_snapshot_digest)?;
170        require_digest("snapshot_digest", &self.snapshot_digest)?;
171        self.tool_results.validate()?;
172        if self.snapshot_digest != self.expected_digest()? {
173            return Err(HarnessEvidenceError::DigestMismatch("snapshot_digest"));
174        }
175        Ok(())
176    }
177
178    pub fn validate_against(
179        &self,
180        input: &ModelInputSnapshotV1,
181    ) -> Result<(), HarnessEvidenceError> {
182        self.validate()?;
183        input.validate()?;
184        if self.call_sequence != input.call_sequence {
185            return Err(HarnessEvidenceError::InvalidContents(
186                "usage and input call sequences agree",
187            ));
188        }
189        if self.input_snapshot_digest != input.snapshot_digest {
190            return Err(HarnessEvidenceError::DigestMismatch(
191                "input_snapshot_digest",
192            ));
193        }
194        if self.estimated_prompt_tokens != input.estimated_prompt_tokens {
195            return Err(HarnessEvidenceError::InvalidContents(
196                "usage and input prompt estimates agree",
197            ));
198        }
199        if self.tool_results.total_count != input.tool_result_count {
200            return Err(HarnessEvidenceError::InvalidContents(
201                "usage and input Tool-result counts agree",
202            ));
203        }
204        Ok(())
205    }
206
207    fn expected_digest(&self) -> Result<String, HarnessEvidenceError> {
208        #[derive(Serialize)]
209        struct Identity<'a> {
210            schema: &'a str,
211            call_sequence: u64,
212            input_snapshot_digest: &'a str,
213            estimated_prompt_tokens: usize,
214            reported_prompt_tokens: usize,
215            reported_completion_tokens: usize,
216            reported_total_tokens: usize,
217            reported_cache_read_tokens: Option<usize>,
218            reported_cache_write_tokens: Option<usize>,
219            tool_results: &'a ToolResultContextUsageV1,
220        }
221
222        Ok(measure(
223            MODEL_USAGE_SNAPSHOT_DOMAIN,
224            &Identity {
225                schema: &self.schema,
226                call_sequence: self.call_sequence,
227                input_snapshot_digest: &self.input_snapshot_digest,
228                estimated_prompt_tokens: self.estimated_prompt_tokens,
229                reported_prompt_tokens: self.reported_prompt_tokens,
230                reported_completion_tokens: self.reported_completion_tokens,
231                reported_total_tokens: self.reported_total_tokens,
232                reported_cache_read_tokens: self.reported_cache_read_tokens,
233                reported_cache_write_tokens: self.reported_cache_write_tokens,
234                tool_results: &self.tool_results,
235            },
236        )?
237        .digest)
238    }
239}
a3s_code_core/harness_evidence/usage.rs

a3s_code_core/harness_evidence/
usage.rs
