poly_kv/
receipt.rs

1use serde::{Deserialize, Serialize};
2
3use crate::digest_compat::Digest;
4use crate::ids_compat::AgentId;
5use crate::policy::CompressionPolicy;
6/// Schema version for receipts.
7pub const RECEIPT_SCHEMA: &str = "poly_kv_receipt_v1";
8/// Schema version for pool build receipts.
9pub const POOL_BUILD_RECEIPT_SCHEMA: &str = "pool_build_receipt_v1";
10/// Schema version for shell materialize receipts.
11pub const SHELL_MATERIALIZE_RECEIPT_SCHEMA: &str = "shell_materialize_receipt_v1";
12/// Schema version for injection receipts.
13pub const INJECTION_RECEIPT_SCHEMA: &str = "injection_receipt_v1";
14/// Schema version for compressed attention selection receipts.
15pub const COMPRESSED_ATTENTION_SELECTION_RECEIPT_SCHEMA: &str =
16    "compressed_attention_selection_receipt_v1";
17
18/// Receipt produced when building a SharedKVPool.
19///
20/// Content-addressed via blake3 digest of canonical JSON.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub struct PoolBuildReceipt {
23    /// Stable schema marker.
24    pub schema_version: String,
25    /// Blake3 digest of the entire pool.
26    pub pool_digest: Digest,
27    /// Per-layer blake3 digests.
28    pub layer_digests: Vec<Digest>,
29    /// Digest of the fib-quant codebook.
30    pub codebook_digest: Digest,
31    /// Digest of the fib-quant rotation.
32    pub rotation_digest: Digest,
33    /// Total number of tokens in the pool.
34    pub total_tokens: u32,
35    /// Milliseconds spent on fib-quant build.
36    pub fib_build_ms: u64,
37    /// Total compressed pool size in bytes.
38    pub pool_size_bytes: u64,
39    /// Total raw f32 size in bytes.
40    pub raw_size_bytes: u64,
41    /// Compression ratio: raw / compressed.
42    pub compression_ratio: f64,
43    /// Snapshot of the policy used.
44    pub policy_snapshot: CompressionPolicy,
45    /// Seed used for construction.
46    pub seeded_with: u64,
47    /// Unix timestamp when built.
48    pub built_at_unix: i64,
49    /// Backend used: "cpu" or "gpu".
50    #[serde(default = "default_backend")]
51    pub backend: String,
52}
53
54fn default_backend() -> String {
55    "cpu".to_string()
56}
57
58impl PoolBuildReceipt {
59    /// Create a new pool build receipt.
60    #[allow(clippy::too_many_arguments)]
61    pub fn new(
62        pool_digest: Digest,
63        layer_digests: Vec<Digest>,
64        codebook_digest: Digest,
65        rotation_digest: Digest,
66        total_tokens: u32,
67        fib_build_ms: u64,
68        pool_size_bytes: u64,
69        raw_size_bytes: u64,
70        policy_snapshot: CompressionPolicy,
71        seeded_with: u64,
72        built_at_unix: i64,
73    ) -> Self {
74        let compression_ratio = if pool_size_bytes > 0 {
75            raw_size_bytes as f64 / pool_size_bytes as f64
76        } else {
77            0.0
78        };
79
80        Self {
81            schema_version: POOL_BUILD_RECEIPT_SCHEMA.into(),
82            pool_digest,
83            layer_digests,
84            codebook_digest,
85            rotation_digest,
86            total_tokens,
87            fib_build_ms,
88            pool_size_bytes,
89            raw_size_bytes,
90            compression_ratio,
91            policy_snapshot,
92            seeded_with,
93            built_at_unix,
94            backend: "cpu".to_string(),
95        }
96    }
97
98    /// Set the backend used for this build.
99    pub fn with_backend(mut self, backend: &str) -> Self {
100        self.backend = backend.to_string();
101        self
102    }
103
104    /// Validates this receipt schema version.
105    pub fn validate(&self) -> crate::error::Result<()> {
106        if self.schema_version != POOL_BUILD_RECEIPT_SCHEMA {
107            return Err(crate::error::PolyKvError::InvalidReceipt(format!(
108                "expected schema {}, got {}",
109                POOL_BUILD_RECEIPT_SCHEMA, self.schema_version
110            )));
111        }
112        if self.pool_digest.hex().is_empty() {
113            return Err(crate::error::PolyKvError::InvalidReceipt(
114                "pool_digest is empty".into(),
115            ));
116        }
117        Ok(())
118    }
119
120    /// Compute the canonical digest of this receipt.
121    pub fn digest(&self) -> crate::error::Result<Digest> {
122        crate::digest_compat::compute_json(self)
123    }
124}
125
126/// Receipt produced when materializing an AgentShell from a pool.
127///
128/// Content-addressed via blake3 digest of canonical JSON.
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
130pub struct ShellMaterializeReceipt {
131    /// Stable schema marker.
132    pub schema_version: String,
133    /// Agent identifier.
134    pub agent_id: AgentId,
135    /// Digest of the parent pool.
136    pub pool_digest: Digest,
137    /// Digest of the materialized shell.
138    pub shell_digest: Digest,
139    /// Number of unique tokens in this shell.
140    pub num_unique_tokens: u32,
141    /// Total compressed shell size in bytes.
142    pub shell_size_bytes: u64,
143    /// Milliseconds spent on materialization.
144    pub materialize_ms: u64,
145    /// Unix timestamp when materialized.
146    pub materialized_at_unix: i64,
147}
148
149impl ShellMaterializeReceipt {
150    /// Create a new shell materialize receipt.
151    pub fn new(
152        agent_id: AgentId,
153        pool_digest: Digest,
154        shell_digest: Digest,
155        num_unique_tokens: u32,
156        shell_size_bytes: u64,
157        materialize_ms: u64,
158        materialized_at_unix: i64,
159    ) -> Self {
160        Self {
161            schema_version: SHELL_MATERIALIZE_RECEIPT_SCHEMA.into(),
162            agent_id,
163            pool_digest,
164            shell_digest,
165            num_unique_tokens,
166            shell_size_bytes,
167            materialize_ms,
168            materialized_at_unix,
169        }
170    }
171
172    /// Validate this receipt.
173    pub fn validate(&self) -> crate::error::Result<()> {
174        if self.schema_version != SHELL_MATERIALIZE_RECEIPT_SCHEMA {
175            return Err(crate::error::PolyKvError::InvalidReceipt(format!(
176                "expected schema {}, got {}",
177                SHELL_MATERIALIZE_RECEIPT_SCHEMA, self.schema_version
178            )));
179        }
180        if self.agent_id.is_empty() {
181            return Err(crate::error::PolyKvError::InvalidReceipt(
182                "agent_id is empty".into(),
183            ));
184        }
185        if self.pool_digest.hex().is_empty() {
186            return Err(crate::error::PolyKvError::InvalidReceipt(
187                "pool_digest is empty".into(),
188            ));
189        }
190        if self.shell_digest.hex().is_empty() {
191            return Err(crate::error::PolyKvError::InvalidReceipt(
192                "shell_digest is empty".into(),
193            ));
194        }
195        Ok(())
196    }
197
198    /// Compute the canonical digest of this receipt.
199    pub fn digest(&self) -> crate::error::Result<Digest> {
200        crate::digest_compat::compute_json(self)
201    }
202}
203
204/// Receipt produced when selecting attention candidates from compressed KV artifacts.
205#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
206pub struct CompressedAttentionSelectionReceipt {
207    /// Stable schema marker.
208    pub schema_version: String,
209    /// Pool digest used for candidate selection.
210    pub pool_digest: Digest,
211    /// Layer index.
212    pub layer: u32,
213    /// KV head index.
214    pub head: u32,
215    /// Number of candidate tokens scored in compressed form.
216    pub candidate_count: u32,
217    /// Number of selected hits returned.
218    pub selected_count: u32,
219    /// Number of cold-pool candidates scored.
220    #[serde(default)]
221    pub pool_candidate_count: u32,
222    /// Number of hot-shell candidates scored.
223    #[serde(default)]
224    pub shell_candidate_count: u32,
225    /// Number of selected hits from the cold pool.
226    #[serde(default)]
227    pub selected_pool_count: u32,
228    /// Number of selected hits from the hot shell.
229    #[serde(default)]
230    pub selected_shell_count: u32,
231    /// Number of compressed key codes scored.
232    pub compressed_key_scores: u64,
233    /// Number of value vectors decoded after top-k selection.
234    pub decoded_value_vectors: u64,
235    /// True only when the full layer was decoded before scoring.
236    pub full_layer_decoded: bool,
237    /// Whether downstream model-quality claims require exact/logit/PPL fallback.
238    #[serde(default = "default_exact_fallback_required")]
239    pub exact_fallback_required: bool,
240    /// Human-readable scoring path / claim boundary.
241    pub scoring_path: String,
242    /// Codec used by the cold shared pool.
243    pub codec_id: String,
244    /// Optional agent id when the selection includes a hot shell.
245    #[serde(default, skip_serializing_if = "Option::is_none")]
246    pub agent_id: Option<String>,
247    /// Optional hot-shell digest when the selection includes a hot shell.
248    #[serde(default, skip_serializing_if = "Option::is_none")]
249    pub shell_digest: Option<Digest>,
250    /// Explicit claim boundary for this selection receipt.
251    #[serde(default = "default_compressed_attention_claim_boundary")]
252    pub claim_boundary: String,
253    /// Unix timestamp when selected.
254    pub selected_at_unix: i64,
255}
256
257fn default_exact_fallback_required() -> bool {
258    true
259}
260
261fn default_compressed_attention_claim_boundary() -> String {
262    "compressed candidate selection only; model-quality/KV-cache preservation claims require exact attention and logit/PPL replay receipts".to_string()
263}
264
265impl CompressedAttentionSelectionReceipt {
266    #[allow(clippy::too_many_arguments)]
267    pub fn new(
268        pool_digest: Digest,
269        layer: u32,
270        head: u32,
271        candidate_count: u32,
272        selected_count: u32,
273        compressed_key_scores: u64,
274        decoded_value_vectors: u64,
275        full_layer_decoded: bool,
276        scoring_path: impl Into<String>,
277        codec_id: impl Into<String>,
278        selected_at_unix: i64,
279    ) -> Self {
280        Self {
281            schema_version: COMPRESSED_ATTENTION_SELECTION_RECEIPT_SCHEMA.into(),
282            pool_digest,
283            layer,
284            head,
285            candidate_count,
286            selected_count,
287            pool_candidate_count: candidate_count,
288            shell_candidate_count: 0,
289            selected_pool_count: selected_count,
290            selected_shell_count: 0,
291            compressed_key_scores,
292            decoded_value_vectors,
293            full_layer_decoded,
294            exact_fallback_required: true,
295            scoring_path: scoring_path.into(),
296            codec_id: codec_id.into(),
297            agent_id: None,
298            shell_digest: None,
299            claim_boundary: default_compressed_attention_claim_boundary(),
300            selected_at_unix,
301        }
302    }
303
304    #[allow(clippy::too_many_arguments)]
305    pub fn with_shell_source_counts(
306        mut self,
307        agent_id: impl Into<String>,
308        shell_digest: Digest,
309        pool_candidate_count: u32,
310        shell_candidate_count: u32,
311        selected_pool_count: u32,
312        selected_shell_count: u32,
313    ) -> Self {
314        self.agent_id = Some(agent_id.into());
315        self.shell_digest = Some(shell_digest);
316        self.pool_candidate_count = pool_candidate_count;
317        self.shell_candidate_count = shell_candidate_count;
318        self.selected_pool_count = selected_pool_count;
319        self.selected_shell_count = selected_shell_count;
320        self
321    }
322
323    pub fn validate(&self) -> crate::error::Result<()> {
324        if self.schema_version != COMPRESSED_ATTENTION_SELECTION_RECEIPT_SCHEMA {
325            return Err(crate::error::PolyKvError::InvalidReceipt(format!(
326                "expected schema {}, got {}",
327                COMPRESSED_ATTENTION_SELECTION_RECEIPT_SCHEMA, self.schema_version
328            )));
329        }
330        if self.pool_digest.hex().is_empty() {
331            return Err(crate::error::PolyKvError::InvalidReceipt(
332                "pool_digest is empty".into(),
333            ));
334        }
335        if self.full_layer_decoded {
336            return Err(crate::error::PolyKvError::InvalidReceipt(
337                "compressed attention selection receipt must not claim full_layer_decoded".into(),
338            ));
339        }
340        if self.selected_count as u64 != self.decoded_value_vectors {
341            return Err(crate::error::PolyKvError::InvalidReceipt(
342                "selected_count must equal decoded_value_vectors for top-k value decode".into(),
343            ));
344        }
345        if self.pool_candidate_count + self.shell_candidate_count != self.candidate_count {
346            return Err(crate::error::PolyKvError::InvalidReceipt(
347                "source candidate counts must sum to candidate_count".into(),
348            ));
349        }
350        if self.selected_pool_count + self.selected_shell_count != self.selected_count {
351            return Err(crate::error::PolyKvError::InvalidReceipt(
352                "selected source counts must sum to selected_count".into(),
353            ));
354        }
355        if self.agent_id.is_some() != self.shell_digest.is_some() {
356            return Err(crate::error::PolyKvError::InvalidReceipt(
357                "agent_id and shell_digest must be present together".into(),
358            ));
359        }
360        if self.claim_boundary.trim().is_empty() {
361            return Err(crate::error::PolyKvError::InvalidReceipt(
362                "claim_boundary is empty".into(),
363            ));
364        }
365        Ok(())
366    }
367
368    pub fn digest(&self) -> crate::error::Result<Digest> {
369        crate::digest_compat::compute_json(self)
370    }
371}
372
373/// Receipt produced when injecting a shell into a KV cache.
374///
375/// Traces the source of every injected block.
376#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
377pub struct InjectionReceipt {
378    /// Stable schema marker.
379    pub schema_version: String,
380    /// Agent identifier.
381    pub agent_id: AgentId,
382    /// Digest of the parent pool.
383    pub pool_digest: Digest,
384    /// Digest of the injected shell.
385    pub shell_digest: Digest,
386    /// Number of injected blocks.
387    pub blocks_injected: u32,
388    /// Per-block digest traces (source -> target).
389    pub block_traces: Vec<BlockInjectionTrace>,
390    /// Unix timestamp when injection occurred.
391    pub injected_at_unix: i64,
392    /// Trace context for cross-boundary correlation (requires typed-ids).
393    #[cfg(feature = "typed-ids")]
394    #[serde(skip_serializing_if = "Option::is_none")]
395    pub trace_ctx: Option<stack_ids::TraceCtx>,
396}
397
398/// Traces one block's injection from source (pool or shell) to target cache.
399#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
400pub struct BlockInjectionTrace {
401    /// Layer index.
402    pub layer: u32,
403    /// Source: "pool" or "shell".
404    pub source: String,
405    /// Digest of the source block.
406    pub source_digest: Digest,
407    /// Position in the target cache.
408    pub target_position: u32,
409}
410
411impl InjectionReceipt {
412    /// Create a new injection receipt.
413    pub fn new(
414        agent_id: AgentId,
415        pool_digest: Digest,
416        shell_digest: Digest,
417        blocks_injected: u32,
418        block_traces: Vec<BlockInjectionTrace>,
419        injected_at_unix: i64,
420    ) -> Self {
421        Self {
422            schema_version: INJECTION_RECEIPT_SCHEMA.into(),
423            agent_id,
424            pool_digest,
425            shell_digest,
426            blocks_injected,
427            block_traces,
428            injected_at_unix,
429            #[cfg(feature = "typed-ids")]
430            trace_ctx: None,
431        }
432    }
433
434    /// Validate this receipt.
435    pub fn validate(&self) -> crate::error::Result<()> {
436        if self.schema_version != INJECTION_RECEIPT_SCHEMA {
437            return Err(crate::error::PolyKvError::InvalidReceipt(format!(
438                "expected schema {}, got {}",
439                INJECTION_RECEIPT_SCHEMA, self.schema_version
440            )));
441        }
442        if self.agent_id.is_empty() {
443            return Err(crate::error::PolyKvError::InvalidReceipt(
444                "agent_id is empty".into(),
445            ));
446        }
447        if self.pool_digest.hex().is_empty() {
448            return Err(crate::error::PolyKvError::InvalidReceipt(
449                "pool_digest is empty".into(),
450            ));
451        }
452        if self.shell_digest.hex().is_empty() {
453            return Err(crate::error::PolyKvError::InvalidReceipt(
454                "shell_digest is empty".into(),
455            ));
456        }
457        Ok(())
458    }
459
460    /// Compute the canonical digest of this receipt.
461    pub fn digest(&self) -> crate::error::Result<Digest> {
462        crate::digest_compat::compute_json(self)
463    }
464
465    /// Builder: set the trace context (requires typed-ids).
466    #[cfg(feature = "typed-ids")]
467    pub fn with_trace_ctx(mut self, ctx: stack_ids::TraceCtx) -> Self {
468        self.trace_ctx = Some(ctx);
469        self
470    }
471}
472
473/// Get the current unix timestamp as i64.
474pub fn now_unix() -> i64 {
475    use std::time::{SystemTime, UNIX_EPOCH};
476    SystemTime::now()
477        .duration_since(UNIX_EPOCH)
478        .unwrap_or_default()
479        .as_secs() as i64
480}
481
482#[cfg(test)]
483mod tests {
484    use super::*;
485    use crate::policy::CompressionPolicy;
486
487    #[test]
488    fn test_pool_build_receipt_round_trip() {
489        let receipt = PoolBuildReceipt::new(
490            Digest::from_hex_unchecked("abc123"),
491            vec![
492                Digest::from_hex_unchecked("layer0_digest"),
493                Digest::from_hex_unchecked("layer1_digest"),
494            ],
495            Digest::from_hex_unchecked("codebook_digest"),
496            Digest::from_hex_unchecked("rotation_digest"),
497            100,
498            42,
499            10_000,
500            500_000,
501            CompressionPolicy::default_two_tier(),
502            42,
503            now_unix(),
504        );
505        assert!(receipt.validate().is_ok());
506
507        let json = serde_json::to_string(&receipt).unwrap();
508        let deser: PoolBuildReceipt = serde_json::from_str(&json).unwrap();
509        assert_eq!(receipt.pool_digest, deser.pool_digest);
510        assert_eq!(receipt.layer_digests, deser.layer_digests);
511        assert_eq!(receipt.compression_ratio, deser.compression_ratio);
512    }
513
514    #[test]
515    fn test_shell_receipt_round_trip() {
516        let receipt = ShellMaterializeReceipt::new(
517            AgentId::new("agent_1"),
518            Digest::from_hex_unchecked("pool_abc"),
519            Digest::from_hex_unchecked("shell_xyz"),
520            50,
521            5_000,
522            10,
523            now_unix(),
524        );
525        assert!(receipt.validate().is_ok());
526
527        let json = serde_json::to_string(&receipt).unwrap();
528        let deser: ShellMaterializeReceipt = serde_json::from_str(&json).unwrap();
529        assert_eq!(receipt.shell_digest, deser.shell_digest);
530    }
531
532    #[test]
533    fn test_injection_receipt_traces() {
534        let traces = vec![
535            BlockInjectionTrace {
536                layer: 0,
537                source: "pool".into(),
538                source_digest: Digest::from_hex_unchecked("abc"),
539                target_position: 0,
540            },
541            BlockInjectionTrace {
542                layer: 0,
543                source: "shell".into(),
544                source_digest: Digest::from_hex_unchecked("def"),
545                target_position: 1,
546            },
547        ];
548        let receipt = InjectionReceipt::new(
549            AgentId::new("agent_1"),
550            Digest::from_hex_unchecked("pool_abc"),
551            Digest::from_hex_unchecked("shell_xyz"),
552            2,
553            traces,
554            now_unix(),
555        );
556        assert!(receipt.validate().is_ok());
557        assert_eq!(receipt.blocks_injected, 2);
558    }
559}
poly_kv/receipt.rs

poly_kv/
receipt.rs
