poly_kv/
manifest.rs

1use serde::{Deserialize, Serialize};
2
3use crate::digest_compat::Digest;
4use crate::ids_compat::AgentId;
5use crate::policy::{CodecId, CompressionPolicy};
6use crate::shape::KvTensorShape;
7
8/// Schema version for PoolManifest.
9pub const POOL_MANIFEST_SCHEMA: &str = "pool_manifest_v1";
10/// Schema version for ShellManifest.
11pub const SHELL_MANIFEST_SCHEMA: &str = "shell_manifest_v1";
12
13/// Manifest describing a built SharedKVPool.
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct PoolManifest {
16    /// Stable schema marker.
17    pub schema_version: String,
18    /// Unique pool identifier (blake3 of canonical JSON).
19    pub pool_id: Digest,
20    /// Logical tensor shape of the underlying model.
21    pub shape: KvTensorShape,
22    /// The compression policy used.
23    pub policy: CompressionPolicy,
24    /// Number of shared tokens in the pool.
25    pub num_shared_tokens: u32,
26    /// Number of transformer layers.
27    pub num_layers: u32,
28    /// Total compressed pool size in bytes.
29    pub pool_size_bytes: u64,
30    /// Codec used for shared pool blocks.
31    pub shared_codec: CodecId,
32    /// Overall compression ratio: raw f32 bytes / compressed bytes.
33    pub compression_ratio: f64,
34    /// Unix timestamp when the pool was built.
35    pub built_at_unix: i64,
36    /// Seed used during pool construction.
37    pub build_seed: u64,
38    /// Scope key for partition-aware pool resolution (requires typed-ids).
39    #[cfg(feature = "typed-ids")]
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub scope: Option<stack_ids::ScopeKey>,
42}
43
44impl PoolManifest {
45    /// Create and validate a pool manifest.
46    #[allow(clippy::too_many_arguments)]
47    pub fn new(
48        pool_id: Digest,
49        shape: KvTensorShape,
50        policy: CompressionPolicy,
51        num_shared_tokens: u32,
52        num_layers: u32,
53        pool_size_bytes: u64,
54        raw_size_bytes: u64,
55        build_seed: u64,
56        built_at_unix: i64,
57    ) -> crate::error::Result<Self> {
58        let compression_ratio = if pool_size_bytes > 0 {
59            raw_size_bytes as f64 / pool_size_bytes as f64
60        } else {
61            0.0
62        };
63
64        let manifest = Self {
65            schema_version: POOL_MANIFEST_SCHEMA.into(),
66            pool_id,
67            shape,
68            policy,
69            num_shared_tokens,
70            num_layers,
71            pool_size_bytes,
72            shared_codec: CODEC_IDENTIFIER_SHARED.into(),
73            compression_ratio,
74            built_at_unix,
75            build_seed,
76            #[cfg(feature = "typed-ids")]
77            scope: None,
78        };
79        manifest.validate()?;
80        Ok(manifest)
81    }
82
83    /// Validate this manifest against the schema.
84    pub fn validate(&self) -> crate::error::Result<()> {
85        if self.schema_version != POOL_MANIFEST_SCHEMA {
86            return Err(crate::error::PolyKvError::InvalidManifest(format!(
87                "expected schema {}, got {}",
88                POOL_MANIFEST_SCHEMA, self.schema_version
89            )));
90        }
91        if self.pool_id.hex().is_empty() {
92            return Err(crate::error::PolyKvError::InvalidManifest(
93                "pool_id is empty".into(),
94            ));
95        }
96        self.shape.validate()?;
97        self.policy.validate()?;
98        Ok(())
99    }
100
101    pub fn digest(&self) -> crate::error::Result<Digest> {
102        crate::digest_compat::compute_json(self)
103    }
104
105    /// Builder: set the scope key (requires typed-ids).
106    #[cfg(feature = "typed-ids")]
107    pub fn with_scope(mut self, scope: stack_ids::ScopeKey) -> Self {
108        self.scope = Some(scope);
109        self
110    }
111}
112
113/// Manifest describing a materialized agent shell.
114#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
115pub struct ShellManifest {
116    /// Stable schema marker.
117    pub schema_version: String,
118    /// Agent identifier.
119    pub agent_id: AgentId,
120    /// Digest of the parent pool.
121    pub pool_digest: Digest,
122    /// Number of unique tokens in this shell.
123    pub num_unique_tokens: u32,
124    /// Number of layers with unique tokens.
125    pub num_unique_layers: u32,
126    /// Total compressed shell size in bytes.
127    pub shell_size_bytes: u64,
128    /// Codec used for shell blocks.
129    pub shell_codec: CodecId,
130    /// Unix timestamp when the shell was materialized.
131    pub materialized_at_unix: i64,
132    /// Seed used during materialization.
133    pub materialize_seed: u64,
134}
135
136impl ShellManifest {
137    /// Create and validate a shell manifest.
138    pub fn new(
139        agent_id: AgentId,
140        pool_digest: Digest,
141        num_unique_tokens: u32,
142        num_unique_layers: u32,
143        shell_size_bytes: u64,
144        materialize_seed: u64,
145        materialized_at_unix: i64,
146    ) -> crate::error::Result<Self> {
147        let manifest = Self {
148            schema_version: SHELL_MANIFEST_SCHEMA.into(),
149            agent_id,
150            pool_digest,
151            num_unique_tokens,
152            num_unique_layers,
153            shell_size_bytes,
154            shell_codec: CODEC_IDENTIFIER_SHELL.into(),
155            materialized_at_unix,
156            materialize_seed,
157        };
158        manifest.validate()?;
159        Ok(manifest)
160    }
161
162    /// Validate this manifest against the schema.
163    pub fn validate(&self) -> crate::error::Result<()> {
164        if self.schema_version != SHELL_MANIFEST_SCHEMA {
165            return Err(crate::error::PolyKvError::InvalidManifest(format!(
166                "expected schema {}, got {}",
167                SHELL_MANIFEST_SCHEMA, self.schema_version
168            )));
169        }
170        if self.agent_id.is_empty() {
171            return Err(crate::error::PolyKvError::InvalidManifest(
172                "agent_id is empty".into(),
173            ));
174        }
175        if self.pool_digest.hex().is_empty() {
176            return Err(crate::error::PolyKvError::InvalidManifest(
177                "pool_digest is empty".into(),
178            ));
179        }
180        Ok(())
181    }
182
183    /// Compute the canonical digest of this manifest.
184    pub fn digest(&self) -> crate::error::Result<Digest> {
185        crate::digest_compat::compute_json(self)
186    }
187}
188
189// Well-known codec identifiers (mirrored in policy.rs, referenced here for manifest usage).
190const CODEC_IDENTIFIER_SHARED: &str = crate::policy::CODEC_FIB_K4_N32;
191const CODEC_IDENTIFIER_SHELL: &str = crate::policy::CODEC_TURBO_8BIT;
poly_kv/manifest.rs

poly_kv/
manifest.rs
