lean_ctx/core/ocla/builtin/
compression_provider.rs1use lean_ctx_protocol::CapabilityManifestV1;
9use std::collections::BTreeMap;
10use std::path::{Path, PathBuf};
11use std::sync::OnceLock;
12
13use crate::core::compressor;
14use crate::core::config::Config;
15use crate::core::ocla::OclaError;
16use crate::core::ocla::content_port::CompressionContentPort;
17use crate::core::ocla::traits::{CompressionProvider, OclaService};
18use crate::core::ocla::types::{
19 CompressionRequest, CompressionResult, OCLA_API_VERSION, OclaCapability, OclaCapabilityKind,
20 OclaCapabilityStatus, OclaResult,
21};
22use crate::core::tokens;
23
24const COMPRESSION_MANIFEST_JSON: &str = include_str!(concat!(
25 env!("CARGO_MANIFEST_DIR"),
26 "/assets/ocla/capability-manifests/leanctx/context-optimization-v1.json"
27));
28
29static COMPRESSION_MANIFEST: OnceLock<Result<CapabilityManifestV1, String>> = OnceLock::new();
30
31static DEFAULT_PORT: OnceLock<Option<CompressionContentPort>> = OnceLock::new();
32
33fn port_from_project_root(root_str: Option<String>) -> Option<CompressionContentPort> {
34 let root_str = root_str?;
35 if root_str.trim().is_empty() || root_str == "." {
36 return None;
37 }
38 let root = PathBuf::from(root_str);
39 let mut component_path = PathBuf::new();
40 for component in root.components() {
41 component_path.push(component.as_os_str());
42 if component_path
43 .symlink_metadata()
44 .ok()?
45 .file_type()
46 .is_symlink()
47 {
48 return None;
49 }
50 }
51 if !root.symlink_metadata().ok()?.is_dir() {
52 return None;
53 }
54 let canonical = root.canonicalize().ok()?;
55 CompressionContentPort::new(canonical)
56}
57
58fn try_default_port() -> Option<&'static CompressionContentPort> {
59 DEFAULT_PORT
60 .get_or_init(|| port_from_project_root(Config::find_project_root()))
61 .as_ref()
62}
63
64pub struct BuiltinCompressionProvider;
65
66impl BuiltinCompressionProvider {
67 pub fn new() -> Self {
68 Self
69 }
70
71 pub fn manifest(&self) -> OclaResult<CapabilityManifestV1> {
73 COMPRESSION_MANIFEST
74 .get_or_init(|| {
75 let manifest =
76 serde_json::from_str::<CapabilityManifestV1>(COMPRESSION_MANIFEST_JSON)
77 .map_err(|error| {
78 format!("invalid pinned compression capability manifest: {error}")
79 })?;
80 manifest.validate().map_err(|error| {
81 format!("invalid pinned compression capability manifest: {error}")
82 })?;
83 Ok(manifest)
84 })
85 .clone()
86 .map_err(OclaError::InvalidRequest)
87 }
88
89 pub fn compress_with_port(
90 &self,
91 request: CompressionRequest,
92 port: &CompressionContentPort,
93 ) -> OclaResult<CompressionResult> {
94 if request.source_tokens == 0 {
95 return Err(OclaError::InvalidRequest(
96 "source_tokens must be > 0".into(),
97 ));
98 }
99 if request.target_tokens == 0 {
100 return Err(OclaError::InvalidRequest(
101 "target_tokens must be > 0".into(),
102 ));
103 }
104
105 if !request.source_ref.starts_with("file:") {
106 return Err(OclaError::InvalidRequest(format!(
107 "only file: refs supported, got: {}",
108 request.source_ref.split(':').next().unwrap_or("unknown")
109 )));
110 }
111
112 let bytes = port.resolve(&request.source_ref)?;
113
114 let source_text = std::str::from_utf8(&bytes)
115 .map_err(|_| OclaError::InvalidRequest("source is not valid UTF-8".into()))?;
116
117 let ext = request
118 .source_ref
119 .strip_prefix("file:")
120 .and_then(|p| Path::new(p).extension())
121 .and_then(|e| e.to_str());
122
123 let compressed = compressor::aggressive_compress(source_text, ext);
124 let delivered_tokens = tokens::count_tokens(&compressed) as u64;
125
126 if delivered_tokens >= request.source_tokens {
127 return Err(OclaError::InvalidRequest(
128 "compression produced no gain (output >= source)".into(),
129 ));
130 }
131
132 if delivered_tokens > request.target_tokens {
133 return Err(OclaError::InvalidRequest(format!(
134 "compressed output ({delivered_tokens}) exceeds target ({})",
135 request.target_tokens
136 )));
137 }
138
139 let ref_key = port.persist(compressed.as_bytes())?;
140
141 Ok(CompressionResult {
142 delivered_ref: ref_key,
143 delivered_tokens,
144 recovery_ref: Some(request.source_ref),
145 })
146 }
147}
148
149impl Default for BuiltinCompressionProvider {
150 fn default() -> Self {
151 Self::new()
152 }
153}
154
155impl OclaService for BuiltinCompressionProvider {
156 fn capability(&self) -> OclaCapability {
157 if try_default_port().is_some() {
158 OclaCapability::available(OclaCapabilityKind::CompressionProvider)
159 } else {
160 OclaCapability {
161 kind: OclaCapabilityKind::CompressionProvider,
162 api_version: OCLA_API_VERSION.to_string(),
163 status: OclaCapabilityStatus::Unavailable,
164 limits: BTreeMap::new(),
165 }
166 }
167 }
168
169 fn manifest(&self) -> OclaResult<CapabilityManifestV1> {
170 BuiltinCompressionProvider::manifest(self)
171 }
172}
173
174impl CompressionProvider for BuiltinCompressionProvider {
175 fn compress(&self, request: CompressionRequest) -> OclaResult<CompressionResult> {
176 let port = try_default_port().ok_or_else(|| {
177 OclaError::InvalidRequest("compression unavailable: no valid project root".into())
178 })?;
179 self.compress_with_port(request, port)
180 }
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186 use crate::core::ocla::types::OclaRequestContext;
187 use std::fs;
188
189 fn ctx() -> OclaRequestContext {
190 OclaRequestContext {
191 request_id: "r1".into(),
192 session_id: "s1".into(),
193 agent_id: "agent-test".into(),
194 content_ref: "blake3:test".into(),
195 tenant_id: None,
196 trace_id: "tr-unit".into(),
197 task_id: None,
198 parent_task_id: None,
199 }
200 }
201
202 #[test]
203 fn compress_rejects_non_file_ref() {
204 let dir = tempfile::tempdir().unwrap();
205 let root = dir.path().canonicalize().unwrap();
206 let port = CompressionContentPort::new(&root).unwrap();
207 let provider = BuiltinCompressionProvider::new();
208 let err = provider
209 .compress_with_port(
210 CompressionRequest {
211 context: ctx(),
212 source_ref: "mem:buffer-123".into(),
213 source_tokens: 1000,
214 target_tokens: 300,
215 quality_policy_ref: None,
216 },
217 &port,
218 )
219 .unwrap_err();
220 assert!(err.to_string().contains("only file: refs"));
221 }
222
223 #[test]
224 fn compress_rejects_zero_source_tokens() {
225 let dir = tempfile::tempdir().unwrap();
226 let root = dir.path().canonicalize().unwrap();
227 let port = CompressionContentPort::new(&root).unwrap();
228 let provider = BuiltinCompressionProvider::new();
229 let err = provider
230 .compress_with_port(
231 CompressionRequest {
232 context: ctx(),
233 source_ref: "file:test.rs".into(),
234 source_tokens: 0,
235 target_tokens: 300,
236 quality_policy_ref: None,
237 },
238 &port,
239 )
240 .unwrap_err();
241 assert!(err.to_string().contains("source_tokens"));
242 }
243
244 #[test]
245 fn compress_rejects_zero_target_tokens() {
246 let dir = tempfile::tempdir().unwrap();
247 let root = dir.path().canonicalize().unwrap();
248 let port = CompressionContentPort::new(&root).unwrap();
249 let provider = BuiltinCompressionProvider::new();
250 let err = provider
251 .compress_with_port(
252 CompressionRequest {
253 context: ctx(),
254 source_ref: "file:test.rs".into(),
255 source_tokens: 100,
256 target_tokens: 0,
257 quality_policy_ref: None,
258 },
259 &port,
260 )
261 .unwrap_err();
262 assert!(err.to_string().contains("target_tokens"));
263 }
264
265 #[test]
266 fn compress_with_real_file_returns_blake3_ref() {
267 let dir = tempfile::tempdir().unwrap();
268 let root = dir.path().canonicalize().unwrap();
269 let content = "use std::collections::HashMap;\n\
270 use std::io::{self, Read, Write, BufReader, BufWriter};\n\
271 use std::fs::File;\n\n\
272 /// Verbose doc that should compress well.\n\
273 /// Another line of documentation.\n\
274 /// Even more verbose documentation here.\n\
275 fn main() -> io::Result<()> {\n \
276 // Initialize the hashmap for storing key-value pairs\n \
277 let mut hash_map_instance: HashMap<String, String> = HashMap::new();\n \
278 // Insert the first key-value pair into the hashmap\n \
279 hash_map_instance.insert(String::from(\"key_one\"), String::from(\"value_one\"));\n \
280 // Insert the second key-value pair into the hashmap\n \
281 hash_map_instance.insert(String::from(\"key_two\"), String::from(\"value_two\"));\n \
282 // Insert the third key-value pair into the hashmap\n \
283 hash_map_instance.insert(String::from(\"key_three\"), String::from(\"value_three\"));\n \
284 // Iterate over all key-value pairs and print them\n \
285 for (key_variable, value_variable) in hash_map_instance.iter() {\n \
286 // Print the current key and value\n \
287 println!(\"Key: {}, Value: {}\", key_variable, value_variable);\n \
288 }\n \
289 // Open a file for reading\n \
290 let input_file_handle: File = File::open(\"input.txt\")?;\n \
291 // Create a buffered reader\n \
292 let mut buffered_reader_instance: BufReader<File> = BufReader::new(input_file_handle);\n \
293 // Read all contents into a string\n \
294 let mut file_contents_buffer: String = String::new();\n \
295 buffered_reader_instance.read_to_string(&mut file_contents_buffer)?;\n \
296 // Print the length\n \
297 println!(\"Read {} bytes from input file\", file_contents_buffer.len());\n \
298 Ok(())\n\
299 }\n";
300 fs::write(root.join("test.rs"), content).unwrap();
301
302 let port = CompressionContentPort::new(&root).unwrap();
303 let provider = BuiltinCompressionProvider::new();
304 let source_tokens = tokens::count_tokens(content) as u64;
305
306 let result = provider.compress_with_port(
307 CompressionRequest {
308 context: ctx(),
309 source_ref: "file:test.rs".into(),
310 source_tokens,
311 target_tokens: source_tokens,
312 quality_policy_ref: None,
313 },
314 &port,
315 );
316
317 let r = result.unwrap();
318 assert!(r.delivered_ref.starts_with("blake3:"));
319 assert!(r.delivered_tokens < source_tokens);
320 assert_eq!(r.recovery_ref, Some("file:test.rs".into()));
321 let delivered = port.retrieve(&r.delivered_ref).unwrap();
322 assert!(delivered.len() < content.len());
323 }
324
325 #[test]
326 fn compress_propagates_resolve_error() {
327 let dir = tempfile::tempdir().unwrap();
328 let root = dir.path().canonicalize().unwrap();
329 let port = CompressionContentPort::new(&root).unwrap();
330 let provider = BuiltinCompressionProvider::new();
331 let err = provider
332 .compress_with_port(
333 CompressionRequest {
334 context: ctx(),
335 source_ref: "file:nonexistent.rs".into(),
336 source_tokens: 100,
337 target_tokens: 50,
338 quality_policy_ref: None,
339 },
340 &port,
341 )
342 .unwrap_err();
343 let msg = err.to_string();
344 assert!(
345 msg.contains("resolve")
346 || msg.contains("No such file")
347 || msg.contains("not found")
348 || msg.contains("cannot find"),
349 "unexpected error: {msg}"
350 );
351 }
352
353 #[test]
354 fn compress_empty_file_persists_empty_output() {
355 let dir = tempfile::tempdir().unwrap();
356 let root = dir.path().canonicalize().unwrap();
357 fs::write(root.join("empty.rs"), b"").unwrap();
358 let port = CompressionContentPort::new(&root).unwrap();
359
360 let result = BuiltinCompressionProvider::new()
361 .compress_with_port(
362 CompressionRequest {
363 context: ctx(),
364 source_ref: "file:empty.rs".into(),
365 source_tokens: 1,
366 target_tokens: 1,
367 quality_policy_ref: None,
368 },
369 &port,
370 )
371 .unwrap();
372
373 assert_eq!(result.delivered_tokens, 0);
374 assert_eq!(result.recovery_ref, Some("file:empty.rs".into()));
375 {
376 let found = &port.retrieve(&result.delivered_ref).unwrap();
377 assert!(found.is_empty(), "{found:?}");
378 }
379 }
380
381 #[test]
382 fn no_gain_rejects_before_persisting() {
383 let dir = tempfile::tempdir().unwrap();
384 let root = dir.path().canonicalize().unwrap();
385 let content = "let value = 1;\n";
386 fs::write(root.join("small.rs"), content).unwrap();
387 let compressed = compressor::aggressive_compress(content, Some("rs"));
388 let source_tokens = tokens::count_tokens(&compressed) as u64;
389 assert!(source_tokens > 0);
390 let ref_key = format!("blake3:{}", blake3::hash(compressed.as_bytes()).to_hex());
391 let port = CompressionContentPort::new(&root).unwrap();
392
393 let err = BuiltinCompressionProvider::new()
394 .compress_with_port(
395 CompressionRequest {
396 context: ctx(),
397 source_ref: "file:small.rs".into(),
398 source_tokens,
399 target_tokens: 100,
400 quality_policy_ref: None,
401 },
402 &port,
403 )
404 .unwrap_err();
405
406 assert!(err.to_string().contains("no gain"));
407 assert!(port.retrieve(&ref_key).is_err());
408 }
409
410 #[test]
411 fn target_cap_rejects_before_persisting() {
412 let dir = tempfile::tempdir().unwrap();
413 let root = dir.path().canonicalize().unwrap();
414 let content = "// repeated documentation\nfn main() {\n let value = 1;\n}\n";
415 fs::write(root.join("target.rs"), content).unwrap();
416 let compressed = compressor::aggressive_compress(content, Some("rs"));
417 let delivered_tokens = tokens::count_tokens(&compressed) as u64;
418 assert!(delivered_tokens > 1);
419 let ref_key = format!("blake3:{}", blake3::hash(compressed.as_bytes()).to_hex());
420 let port = CompressionContentPort::new(&root).unwrap();
421
422 let err = BuiltinCompressionProvider::new()
423 .compress_with_port(
424 CompressionRequest {
425 context: ctx(),
426 source_ref: "file:target.rs".into(),
427 source_tokens: tokens::count_tokens(content) as u64,
428 target_tokens: delivered_tokens - 1,
429 quality_policy_ref: None,
430 },
431 &port,
432 )
433 .unwrap_err();
434
435 assert!(err.to_string().contains("exceeds target"));
436 assert!(port.retrieve(&ref_key).is_err());
437 }
438
439 #[cfg(unix)]
440 #[test]
441 fn project_root_symlink_is_unavailable() {
442 let dir = tempfile::tempdir().unwrap();
443 let link = dir.path().join("root-link");
444 std::os::unix::fs::symlink(dir.path(), &link).unwrap();
445 assert!(port_from_project_root(Some(link.to_string_lossy().into_owned())).is_none());
446 }
447
448 #[test]
449 fn project_root_empty_or_dot_is_unavailable() {
450 assert!(port_from_project_root(Some(String::new())).is_none());
451 assert!(port_from_project_root(Some(".".into())).is_none());
452 }
453
454 #[test]
455 fn capability_reflects_root_availability() {
456 let provider = BuiltinCompressionProvider::new();
457 let cap = provider.capability();
458 assert_eq!(cap.kind, OclaCapabilityKind::CompressionProvider);
459 }
460}