a3s_code_core/context/
mod.rs

1//! Context Provider Extension Point
2//!
3//! This module provides the extension point for integrating context databases
4//! like OpenViking into the agent loop. Context providers can supply memory,
5//! resources, and skills to augment the LLM's context.
6//!
7//! ## Usage
8//!
9//! Implement the `ContextProvider` trait and register it with a session:
10//!
11//! ```ignore
12//! use a3s_code::context::{ContextProvider, ContextQuery, ContextResult};
13//!
14//! struct MyProvider { /* ... */ }
15//!
16//! #[async_trait::async_trait]
17//! impl ContextProvider for MyProvider {
18//!     fn name(&self) -> &str { "my-provider" }
19//!
20//!     async fn query(&self, query: &ContextQuery) -> anyhow::Result<ContextResult> {
21//!         // Retrieve relevant context...
22//!     }
23//! }
24//! ```
25
26pub mod assembler;
27pub mod fs_provider;
28pub mod recent_workspace_provider;
29pub mod ripgrep_provider;
30pub mod skill_catalog_provider;
31pub mod static_provider;
32
33pub use assembler::{
34    ContextAssembler, ContextAssembly, ContextAssemblyPolicy, ContextBudget, ContextSourcePolicy,
35};
36pub use fs_provider::{FileSystemContextConfig, FileSystemContextProvider};
37pub use recent_workspace_provider::RecentWorkspaceFilesContextProvider;
38pub use ripgrep_provider::{RipgrepContextConfig, RipgrepContextProvider};
39pub use skill_catalog_provider::SkillCatalogContextProvider;
40pub use static_provider::StaticContextProvider;
41
42use serde::{Deserialize, Serialize};
43use std::collections::HashMap;
44
45pub const CONTEXT_PROVENANCE_METADATA: &str = "a3s.context.provenance";
46pub const CONTEXT_PRIORITY_METADATA: &str = "a3s.context.priority";
47pub const CONTEXT_TRUST_METADATA: &str = "a3s.context.trust";
48pub const CONTEXT_FRESHNESS_METADATA: &str = "a3s.context.freshness";
49pub(crate) const CONTEXT_REQUIRED_METADATA: &str = "a3s.context.required";
50
51/// Type of context being queried
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
53pub enum ContextType {
54    /// Session/user history, extracted insights
55    Memory,
56    /// Documentation, code, knowledge base
57    #[default]
58    Resource,
59    /// Agent capabilities, behavior instructions
60    Skill,
61}
62
63/// Retrieval depth for tiered context (L0/L1/L2 pattern)
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
65pub enum ContextDepth {
66    /// ~100 tokens - high-level summary
67    Abstract,
68    /// ~2k tokens - key details (default)
69    #[default]
70    Overview,
71    /// Variable - complete content
72    Full,
73}
74
75/// Query to a context provider
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct ContextQuery {
78    /// The query string to search for relevant context
79    pub query: String,
80
81    /// Types of context to retrieve
82    #[serde(default)]
83    pub context_types: Vec<ContextType>,
84
85    /// Desired retrieval depth
86    #[serde(default)]
87    pub depth: ContextDepth,
88
89    /// Maximum number of results to return
90    #[serde(default = "default_max_results")]
91    pub max_results: usize,
92
93    /// Maximum total tokens across all results
94    #[serde(default = "default_max_tokens")]
95    pub max_tokens: usize,
96
97    /// Optional session ID for session-specific context
98    #[serde(default)]
99    pub session_id: Option<String>,
100
101    /// Additional provider-specific parameters
102    #[serde(default)]
103    pub params: HashMap<String, serde_json::Value>,
104}
105
106fn default_max_results() -> usize {
107    10
108}
109
110fn default_max_tokens() -> usize {
111    4000
112}
113
114impl ContextQuery {
115    /// Create a new context query with defaults
116    pub fn new(query: impl Into<String>) -> Self {
117        Self {
118            query: query.into(),
119            context_types: vec![ContextType::Resource],
120            depth: ContextDepth::default(),
121            max_results: default_max_results(),
122            max_tokens: default_max_tokens(),
123            session_id: None,
124            params: HashMap::new(),
125        }
126    }
127
128    /// Set the context types to retrieve
129    pub fn with_types(mut self, types: impl IntoIterator<Item = ContextType>) -> Self {
130        self.context_types = types.into_iter().collect();
131        self
132    }
133
134    /// Set the retrieval depth
135    pub fn with_depth(mut self, depth: ContextDepth) -> Self {
136        self.depth = depth;
137        self
138    }
139
140    /// Set the maximum number of results
141    pub fn with_max_results(mut self, max: usize) -> Self {
142        self.max_results = max;
143        self
144    }
145
146    /// Set the maximum total tokens
147    pub fn with_max_tokens(mut self, max: usize) -> Self {
148        self.max_tokens = max;
149        self
150    }
151
152    /// Set the session ID
153    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
154        self.session_id = Some(id.into());
155        self
156    }
157
158    /// Add a custom parameter
159    pub fn with_param(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
160        self.params.insert(key.into(), value);
161        self
162    }
163}
164
165/// A single piece of retrieved context
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct ContextItem {
168    /// Unique identifier for this context item
169    pub id: String,
170
171    /// Type of context
172    pub context_type: ContextType,
173
174    /// The actual content
175    pub content: String,
176
177    /// Estimated token count (informational)
178    #[serde(default)]
179    pub token_count: usize,
180
181    /// Relevance score (0.0 to 1.0)
182    #[serde(default)]
183    pub relevance: f32,
184
185    /// Optional source URI (e.g., "viking://docs/auth")
186    #[serde(default)]
187    pub source: Option<String>,
188
189    /// Additional metadata
190    #[serde(default)]
191    pub metadata: HashMap<String, serde_json::Value>,
192}
193
194impl ContextItem {
195    /// Create a new context item
196    pub fn new(
197        id: impl Into<String>,
198        context_type: ContextType,
199        content: impl Into<String>,
200    ) -> Self {
201        Self {
202            id: id.into(),
203            context_type,
204            content: content.into(),
205            token_count: 0,
206            relevance: 0.0,
207            source: None,
208            metadata: HashMap::new(),
209        }
210    }
211
212    /// Set the token count
213    pub fn with_token_count(mut self, count: usize) -> Self {
214        self.token_count = count;
215        self
216    }
217
218    /// Set the relevance score
219    pub fn with_relevance(mut self, score: f32) -> Self {
220        self.relevance = score.clamp(0.0, 1.0);
221        self
222    }
223
224    /// Set the source URI
225    pub fn with_source(mut self, source: impl Into<String>) -> Self {
226        self.source = Some(source.into());
227        self
228    }
229
230    /// Add metadata
231    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
232        self.metadata.insert(key.into(), value);
233        self
234    }
235
236    /// Set a human-readable provenance label for diagnostics and ranking.
237    pub fn with_provenance(mut self, provenance: impl Into<String>) -> Self {
238        self.metadata.insert(
239            CONTEXT_PROVENANCE_METADATA.to_string(),
240            serde_json::Value::String(provenance.into()),
241        );
242        self
243    }
244
245    /// Set priority score (0.0 to 1.0) for harness-controlled ranking.
246    pub fn with_priority(mut self, priority: f32) -> Self {
247        self.metadata.insert(
248            CONTEXT_PRIORITY_METADATA.to_string(),
249            serde_json::json!(priority.clamp(0.0, 1.0)),
250        );
251        self
252    }
253
254    /// Set trust score (0.0 to 1.0) for harness-controlled ranking.
255    pub fn with_trust(mut self, trust: f32) -> Self {
256        self.metadata.insert(
257            CONTEXT_TRUST_METADATA.to_string(),
258            serde_json::json!(trust.clamp(0.0, 1.0)),
259        );
260        self
261    }
262
263    /// Set freshness score (0.0 to 1.0) for harness-controlled ranking.
264    pub fn with_freshness(mut self, freshness: f32) -> Self {
265        self.metadata.insert(
266            CONTEXT_FRESHNESS_METADATA.to_string(),
267            serde_json::json!(freshness.clamp(0.0, 1.0)),
268        );
269        self
270    }
271
272    /// Mark harness-owned context as mandatory prompt material.
273    ///
274    /// This is reserved for bounded instruction sources such as the effective
275    /// AGENTS.md chain. Ordinary retrieval providers remain subject to the
276    /// normal per-query and assembly budgets.
277    pub(crate) fn with_required(mut self) -> Self {
278        self.metadata.insert(
279            CONTEXT_REQUIRED_METADATA.to_string(),
280            serde_json::Value::Bool(true),
281        );
282        self
283    }
284
285    pub fn provenance(&self) -> Option<&str> {
286        self.metadata
287            .get(CONTEXT_PROVENANCE_METADATA)
288            .and_then(serde_json::Value::as_str)
289    }
290
291    pub fn priority(&self) -> f32 {
292        metadata_score(self.metadata.get(CONTEXT_PRIORITY_METADATA))
293    }
294
295    pub fn trust(&self) -> f32 {
296        metadata_score(self.metadata.get(CONTEXT_TRUST_METADATA))
297    }
298
299    pub fn freshness(&self) -> f32 {
300        metadata_score(self.metadata.get(CONTEXT_FRESHNESS_METADATA))
301    }
302
303    pub(crate) fn is_required(&self) -> bool {
304        self.metadata
305            .get(CONTEXT_REQUIRED_METADATA)
306            .and_then(serde_json::Value::as_bool)
307            .unwrap_or(false)
308    }
309
310    /// Format as XML tag for system prompt injection
311    pub fn to_xml(&self) -> String {
312        let source_attr = self
313            .source
314            .as_ref()
315            .map(|s| format!(" source=\"{}\"", escape_xml_attribute(s)))
316            .unwrap_or_default();
317        let type_str = match self.context_type {
318            ContextType::Memory => "Memory",
319            ContextType::Resource => "Resource",
320            ContextType::Skill => "Skill",
321        };
322        format!(
323            "<context{} type=\"{}\">\n{}\n</context>",
324            source_attr,
325            type_str,
326            escape_xml_text(&self.content)
327        )
328    }
329}
330
331fn escape_xml_attribute(value: &str) -> String {
332    value.chars().fold(String::new(), |mut escaped, character| {
333        match character {
334            '&' => escaped.push_str("&amp;"),
335            '<' => escaped.push_str("&lt;"),
336            '>' => escaped.push_str("&gt;"),
337            '"' => escaped.push_str("&quot;"),
338            '\'' => escaped.push_str("&apos;"),
339            _ => escaped.push(character),
340        }
341        escaped
342    })
343}
344
345fn escape_xml_text(value: &str) -> String {
346    value.chars().fold(String::new(), |mut escaped, character| {
347        match character {
348            '&' => escaped.push_str("&amp;"),
349            '<' => escaped.push_str("&lt;"),
350            '>' => escaped.push_str("&gt;"),
351            _ => escaped.push(character),
352        }
353        escaped
354    })
355}
356
357fn metadata_score(value: Option<&serde_json::Value>) -> f32 {
358    value
359        .and_then(serde_json::Value::as_f64)
360        .map(|score| (score as f32).clamp(0.0, 1.0))
361        .unwrap_or(0.0)
362}
363
364/// Read a UTF-8 file without allowing a concurrent file growth race to bypass
365/// the configured byte limit.
366///
367/// A file that is too large or is not valid UTF-8 is treated as an unsupported
368/// context candidate and returns `Ok(None)`. Actual I/O failures remain errors
369/// so callers can preserve their existing diagnostics.
370pub(crate) fn read_utf8_file_bounded(
371    path: &std::path::Path,
372    max_bytes: usize,
373) -> std::io::Result<Option<String>> {
374    let bytes = match crate::bounded_io::read_file_bounded(path, max_bytes) {
375        Ok(bytes) => bytes,
376        Err(error) if error.kind() == std::io::ErrorKind::InvalidData => return Ok(None),
377        Err(error) => return Err(error),
378    };
379
380    Ok(String::from_utf8(bytes).ok())
381}
382
383/// Result from a context provider query
384#[derive(Debug, Clone, Default, Serialize, Deserialize)]
385pub struct ContextResult {
386    /// Retrieved context items
387    pub items: Vec<ContextItem>,
388
389    /// Total tokens across all items
390    pub total_tokens: usize,
391
392    /// Name of the provider that returned these results
393    pub provider: String,
394
395    /// Whether results were truncated due to limits
396    pub truncated: bool,
397}
398
399/// Runtime behavior when one context provider returns an error.
400#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
401pub enum ContextProviderFailureMode {
402    /// Preserve the existing general-purpose RAG behavior: log and omit the
403    /// failed provider while other sources may continue.
404    #[default]
405    BestEffort,
406    /// Abort the turn. Used by exact-generation cognitive packages so a
407    /// provider failure cannot silently fall back to memory, graph data, or an
408    /// unpinned package.
409    FailClosed,
410}
411
412impl ContextResult {
413    /// Create a new empty result
414    pub fn new(provider: impl Into<String>) -> Self {
415        Self {
416            items: Vec::new(),
417            total_tokens: 0,
418            provider: provider.into(),
419            truncated: false,
420        }
421    }
422
423    /// Add an item to the result
424    pub fn add_item(&mut self, item: ContextItem) {
425        self.total_tokens = self.total_tokens.saturating_add(item.token_count);
426        self.items.push(item);
427    }
428
429    /// Check if the result is empty
430    pub fn is_empty(&self) -> bool {
431        self.items.is_empty()
432    }
433
434    /// Format all items as XML for system prompt injection
435    pub fn to_xml(&self) -> String {
436        self.items
437            .iter()
438            .map(|item| item.to_xml())
439            .collect::<Vec<_>>()
440            .join("\n\n")
441    }
442}
443
444/// Context provider trait - implement this for OpenViking, RAG systems, etc.
445#[async_trait::async_trait]
446pub trait ContextProvider: Send + Sync {
447    /// Provider name (used for identification and logging)
448    fn name(&self) -> &str;
449
450    /// Whether provider failure may be omitted or must abort the turn.
451    fn failure_mode(&self) -> ContextProviderFailureMode {
452        ContextProviderFailureMode::BestEffort
453    }
454
455    /// Exact cognitive-package identity, when this is Code's typed adapter.
456    ///
457    /// Hosts must install such an adapter through
458    /// [`SessionOptions::with_cognitive_context`](crate::SessionOptions::with_cognitive_context)
459    /// so the same value is persisted in the session snapshot.
460    fn cognitive_package_binding(
461        &self,
462    ) -> Option<&crate::cognitive_context::CognitivePackageBindingV1> {
463        None
464    }
465
466    /// Query the provider for relevant context
467    async fn query(&self, query: &ContextQuery) -> anyhow::Result<ContextResult>;
468
469    /// Called after each turn for memory extraction (optional)
470    ///
471    /// Providers can implement this to extract and store memories from
472    /// the conversation.
473    async fn on_turn_complete(
474        &self,
475        _session_id: &str,
476        _prompt: &str,
477        _response: &str,
478    ) -> anyhow::Result<()> {
479        Ok(())
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use super::*;
486
487    // ========================================================================
488    // ContextType Tests
489    // ========================================================================
490
491    #[test]
492    fn test_context_type_default() {
493        let ct: ContextType = Default::default();
494        assert_eq!(ct, ContextType::Resource);
495    }
496
497    #[test]
498    fn test_context_type_serialization() {
499        let ct = ContextType::Memory;
500        let json = serde_json::to_string(&ct).unwrap();
501        assert_eq!(json, "\"Memory\"");
502
503        let parsed: ContextType = serde_json::from_str(&json).unwrap();
504        assert_eq!(parsed, ContextType::Memory);
505    }
506
507    #[test]
508    fn test_context_type_all_variants() {
509        let types = vec![
510            ContextType::Memory,
511            ContextType::Resource,
512            ContextType::Skill,
513        ];
514        for ct in types {
515            let json = serde_json::to_string(&ct).unwrap();
516            let parsed: ContextType = serde_json::from_str(&json).unwrap();
517            assert_eq!(parsed, ct);
518        }
519    }
520
521    // ========================================================================
522    // ContextDepth Tests
523    // ========================================================================
524
525    #[test]
526    fn test_context_depth_default() {
527        let cd: ContextDepth = Default::default();
528        assert_eq!(cd, ContextDepth::Overview);
529    }
530
531    #[test]
532    fn test_context_depth_serialization() {
533        let cd = ContextDepth::Full;
534        let json = serde_json::to_string(&cd).unwrap();
535        assert_eq!(json, "\"Full\"");
536
537        let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
538        assert_eq!(parsed, ContextDepth::Full);
539    }
540
541    #[test]
542    fn test_context_depth_all_variants() {
543        let depths = vec![
544            ContextDepth::Abstract,
545            ContextDepth::Overview,
546            ContextDepth::Full,
547        ];
548        for cd in depths {
549            let json = serde_json::to_string(&cd).unwrap();
550            let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
551            assert_eq!(parsed, cd);
552        }
553    }
554
555    // ========================================================================
556    // ContextQuery Tests
557    // ========================================================================
558
559    #[test]
560    fn test_context_query_new() {
561        let query = ContextQuery::new("test query");
562        assert_eq!(query.query, "test query");
563        assert_eq!(query.context_types, vec![ContextType::Resource]);
564        assert_eq!(query.depth, ContextDepth::Overview);
565        assert_eq!(query.max_results, 10);
566        assert_eq!(query.max_tokens, 4000);
567        assert!(query.session_id.is_none());
568        assert!(query.params.is_empty());
569    }
570
571    #[test]
572    fn test_context_query_builder() {
573        let query = ContextQuery::new("test")
574            .with_types([ContextType::Memory, ContextType::Skill])
575            .with_depth(ContextDepth::Full)
576            .with_max_results(5)
577            .with_max_tokens(2000)
578            .with_session_id("sess-123")
579            .with_param("custom", serde_json::json!("value"));
580
581        assert_eq!(query.context_types.len(), 2);
582        assert!(query.context_types.contains(&ContextType::Memory));
583        assert!(query.context_types.contains(&ContextType::Skill));
584        assert_eq!(query.depth, ContextDepth::Full);
585        assert_eq!(query.max_results, 5);
586        assert_eq!(query.max_tokens, 2000);
587        assert_eq!(query.session_id, Some("sess-123".to_string()));
588        assert_eq!(
589            query.params.get("custom"),
590            Some(&serde_json::json!("value"))
591        );
592    }
593
594    #[test]
595    fn test_context_query_serialization() {
596        let query = ContextQuery::new("search term")
597            .with_types([ContextType::Resource])
598            .with_session_id("sess-456");
599
600        let json = serde_json::to_string(&query).unwrap();
601        let parsed: ContextQuery = serde_json::from_str(&json).unwrap();
602
603        assert_eq!(parsed.query, "search term");
604        assert_eq!(parsed.session_id, Some("sess-456".to_string()));
605    }
606
607    #[test]
608    fn test_context_query_deserialization_with_defaults() {
609        let json = r#"{"query": "minimal query"}"#;
610        let query: ContextQuery = serde_json::from_str(json).unwrap();
611
612        assert_eq!(query.query, "minimal query");
613        assert!(query.context_types.is_empty()); // Default from serde is empty vec
614        assert_eq!(query.depth, ContextDepth::Overview);
615        assert_eq!(query.max_results, 10);
616        assert_eq!(query.max_tokens, 4000);
617    }
618
619    // ========================================================================
620    // ContextItem Tests
621    // ========================================================================
622
623    #[test]
624    fn test_context_item_new() {
625        let item = ContextItem::new("item-1", ContextType::Resource, "Some content");
626        assert_eq!(item.id, "item-1");
627        assert_eq!(item.context_type, ContextType::Resource);
628        assert_eq!(item.content, "Some content");
629        assert_eq!(item.token_count, 0);
630        assert_eq!(item.relevance, 0.0);
631        assert!(item.source.is_none());
632        assert!(item.metadata.is_empty());
633    }
634
635    #[test]
636    fn test_context_item_builder() {
637        let item = ContextItem::new("item-2", ContextType::Memory, "Memory content")
638            .with_token_count(150)
639            .with_relevance(0.85)
640            .with_source("viking://memory/session-123")
641            .with_provenance("memory")
642            .with_priority(0.7)
643            .with_trust(1.2)
644            .with_freshness(-1.0)
645            .with_metadata("key", serde_json::json!("value"));
646
647        assert_eq!(item.token_count, 150);
648        assert!((item.relevance - 0.85).abs() < f32::EPSILON);
649        assert_eq!(item.source, Some("viking://memory/session-123".to_string()));
650        assert_eq!(item.provenance(), Some("memory"));
651        assert!((item.priority() - 0.7).abs() < f32::EPSILON);
652        assert!((item.trust() - 1.0).abs() < f32::EPSILON);
653        assert!(item.freshness().abs() < f32::EPSILON);
654        assert_eq!(item.metadata.get("key"), Some(&serde_json::json!("value")));
655    }
656
657    #[test]
658    fn test_context_item_relevance_clamping() {
659        let item1 = ContextItem::new("id", ContextType::Resource, "").with_relevance(1.5);
660        assert!((item1.relevance - 1.0).abs() < f32::EPSILON);
661
662        let item2 = ContextItem::new("id", ContextType::Resource, "").with_relevance(-0.5);
663        assert!(item2.relevance.abs() < f32::EPSILON);
664    }
665
666    #[test]
667    fn test_context_item_to_xml_without_source() {
668        let item = ContextItem::new("id", ContextType::Resource, "Content here");
669        let xml = item.to_xml();
670        assert_eq!(xml, "<context type=\"Resource\">\nContent here\n</context>");
671    }
672
673    #[test]
674    fn test_context_item_to_xml_with_source() {
675        let item = ContextItem::new("id", ContextType::Memory, "Memory content")
676            .with_source("viking://docs/auth");
677        let xml = item.to_xml();
678        assert_eq!(
679            xml,
680            "<context source=\"viking://docs/auth\" type=\"Memory\">\nMemory content\n</context>"
681        );
682    }
683
684    #[test]
685    fn test_context_item_to_xml_escapes_untrusted_source_and_content() {
686        let item = ContextItem::new(
687            "item",
688            ContextType::Resource,
689            "<directive>& keep \"quotes\"",
690        )
691        .with_source("provider://a\"&b<c>");
692
693        assert_eq!(
694            item.to_xml(),
695            "<context source=\"provider://a&quot;&amp;b&lt;c&gt;\" type=\"Resource\">\n&lt;directive&gt;&amp; keep \"quotes\"\n</context>"
696        );
697    }
698
699    #[test]
700    fn test_context_item_to_xml_all_types() {
701        let memory = ContextItem::new("m", ContextType::Memory, "m").to_xml();
702        assert!(memory.contains("type=\"Memory\""));
703
704        let resource = ContextItem::new("r", ContextType::Resource, "r").to_xml();
705        assert!(resource.contains("type=\"Resource\""));
706
707        let skill = ContextItem::new("s", ContextType::Skill, "s").to_xml();
708        assert!(skill.contains("type=\"Skill\""));
709    }
710
711    #[test]
712    fn test_context_item_serialization() {
713        let item = ContextItem::new("item-3", ContextType::Skill, "Skill instructions")
714            .with_token_count(200)
715            .with_relevance(0.9)
716            .with_source("viking://skills/code-review");
717
718        let json = serde_json::to_string(&item).unwrap();
719        let parsed: ContextItem = serde_json::from_str(&json).unwrap();
720
721        assert_eq!(parsed.id, "item-3");
722        assert_eq!(parsed.context_type, ContextType::Skill);
723        assert_eq!(parsed.content, "Skill instructions");
724        assert_eq!(parsed.token_count, 200);
725    }
726
727    // ========================================================================
728    // ContextResult Tests
729    // ========================================================================
730
731    #[test]
732    fn test_context_result_new() {
733        let result = ContextResult::new("test-provider");
734        assert!(result.items.is_empty());
735        assert_eq!(result.total_tokens, 0);
736        assert_eq!(result.provider, "test-provider");
737        assert!(!result.truncated);
738    }
739
740    #[test]
741    fn test_context_result_add_item() {
742        let mut result = ContextResult::new("provider");
743        let item = ContextItem::new("id", ContextType::Resource, "content").with_token_count(100);
744        result.add_item(item);
745
746        assert_eq!(result.items.len(), 1);
747        assert_eq!(result.total_tokens, 100);
748    }
749
750    #[test]
751    fn test_context_result_add_multiple_items() {
752        let mut result = ContextResult::new("provider");
753        result.add_item(ContextItem::new("1", ContextType::Resource, "a").with_token_count(50));
754        result.add_item(ContextItem::new("2", ContextType::Memory, "b").with_token_count(75));
755        result.add_item(ContextItem::new("3", ContextType::Skill, "c").with_token_count(25));
756
757        assert_eq!(result.items.len(), 3);
758        assert_eq!(result.total_tokens, 150);
759    }
760
761    #[test]
762    fn context_result_token_accounting_saturates_on_overflow() {
763        let mut result = ContextResult::new("provider");
764        result.add_item(
765            ContextItem::new("large", ContextType::Resource, "large").with_token_count(usize::MAX),
766        );
767        result
768            .add_item(ContextItem::new("next", ContextType::Resource, "next").with_token_count(1));
769
770        assert_eq!(result.items.len(), 2);
771        assert_eq!(result.total_tokens, usize::MAX);
772    }
773
774    #[test]
775    fn bounded_context_file_read_rejects_growth_and_invalid_utf8() {
776        use std::io::{Seek, SeekFrom, Write};
777
778        let mut file = tempfile::NamedTempFile::new().unwrap();
779        write!(file, "12345").unwrap();
780        assert!(read_utf8_file_bounded(file.path(), 4).unwrap().is_none());
781        assert_eq!(
782            read_utf8_file_bounded(file.path(), 5).unwrap().as_deref(),
783            Some("12345")
784        );
785
786        file.as_file_mut().set_len(2).unwrap();
787        file.as_file_mut().seek(SeekFrom::Start(0)).unwrap();
788        file.as_file_mut().write_all(&[0xff, 0xfe]).unwrap();
789        assert!(read_utf8_file_bounded(file.path(), 2).unwrap().is_none());
790    }
791
792    #[test]
793    fn test_context_result_is_empty() {
794        let empty = ContextResult::new("provider");
795        assert!(empty.is_empty());
796
797        let mut non_empty = ContextResult::new("provider");
798        non_empty.add_item(ContextItem::new("id", ContextType::Resource, "content"));
799        assert!(!non_empty.is_empty());
800    }
801
802    #[test]
803    fn test_context_result_to_xml() {
804        let mut result = ContextResult::new("provider");
805        result.add_item(
806            ContextItem::new("1", ContextType::Resource, "First content").with_source("source://1"),
807        );
808        result.add_item(ContextItem::new("2", ContextType::Memory, "Second content"));
809
810        let xml = result.to_xml();
811        assert!(xml.contains("<context source=\"source://1\" type=\"Resource\">"));
812        assert!(xml.contains("First content"));
813        assert!(xml.contains("<context type=\"Memory\">"));
814        assert!(xml.contains("Second content"));
815    }
816
817    #[test]
818    fn test_context_result_to_xml_empty() {
819        let result = ContextResult::new("provider");
820        let xml = result.to_xml();
821        assert!(xml.is_empty());
822    }
823
824    #[test]
825    fn test_context_result_serialization() {
826        let mut result = ContextResult::new("test-provider");
827        result.truncated = true;
828        result.add_item(ContextItem::new("id", ContextType::Resource, "content"));
829
830        let json = serde_json::to_string(&result).unwrap();
831        let parsed: ContextResult = serde_json::from_str(&json).unwrap();
832
833        assert_eq!(parsed.provider, "test-provider");
834        assert!(parsed.truncated);
835        assert_eq!(parsed.items.len(), 1);
836    }
837
838    #[test]
839    fn test_context_result_default() {
840        let result: ContextResult = Default::default();
841        assert!(result.items.is_empty());
842        assert_eq!(result.total_tokens, 0);
843        assert!(result.provider.is_empty());
844        assert!(!result.truncated);
845    }
846
847    // ========================================================================
848    // ContextProvider Trait Tests (with Mock)
849    // ========================================================================
850
851    struct MockContextProvider {
852        name: String,
853        items: Vec<ContextItem>,
854    }
855
856    impl MockContextProvider {
857        fn new(name: &str) -> Self {
858            Self {
859                name: name.to_string(),
860                items: Vec::new(),
861            }
862        }
863
864        fn with_items(mut self, items: Vec<ContextItem>) -> Self {
865            self.items = items;
866            self
867        }
868    }
869
870    #[async_trait::async_trait]
871    impl ContextProvider for MockContextProvider {
872        fn name(&self) -> &str {
873            &self.name
874        }
875
876        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
877            let mut result = ContextResult::new(&self.name);
878            for item in &self.items {
879                result.add_item(item.clone());
880            }
881            Ok(result)
882        }
883    }
884
885    #[tokio::test]
886    async fn test_mock_context_provider() {
887        let provider = MockContextProvider::new("mock").with_items(vec![ContextItem::new(
888            "1",
889            ContextType::Resource,
890            "content",
891        )]);
892
893        assert_eq!(provider.name(), "mock");
894
895        let query = ContextQuery::new("test");
896        let result = provider.query(&query).await.unwrap();
897
898        assert_eq!(result.provider, "mock");
899        assert_eq!(result.items.len(), 1);
900    }
901
902    #[tokio::test]
903    async fn test_context_provider_on_turn_complete_default() {
904        let provider = MockContextProvider::new("mock");
905
906        // Default implementation should succeed
907        let result = provider
908            .on_turn_complete("session-1", "prompt", "response")
909            .await;
910        assert!(result.is_ok());
911    }
912
913    struct MockMemoryProvider {
914        memories: std::sync::Arc<tokio::sync::RwLock<Vec<(String, String, String)>>>,
915    }
916
917    impl MockMemoryProvider {
918        fn new() -> Self {
919            Self {
920                memories: std::sync::Arc::new(tokio::sync::RwLock::new(Vec::new())),
921            }
922        }
923    }
924
925    #[async_trait::async_trait]
926    impl ContextProvider for MockMemoryProvider {
927        fn name(&self) -> &str {
928            "memory-provider"
929        }
930
931        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
932            Ok(ContextResult::new("memory-provider"))
933        }
934
935        async fn on_turn_complete(
936            &self,
937            session_id: &str,
938            prompt: &str,
939            response: &str,
940        ) -> anyhow::Result<()> {
941            let mut memories = self.memories.write().await;
942            memories.push((
943                session_id.to_string(),
944                prompt.to_string(),
945                response.to_string(),
946            ));
947            Ok(())
948        }
949    }
950
951    #[tokio::test]
952    async fn test_context_provider_on_turn_complete_custom() {
953        let provider = MockMemoryProvider::new();
954
955        provider
956            .on_turn_complete("sess-1", "What is Rust?", "Rust is a systems language.")
957            .await
958            .unwrap();
959
960        let memories = provider.memories.read().await;
961        assert_eq!(memories.len(), 1);
962        assert_eq!(memories[0].0, "sess-1");
963        assert_eq!(memories[0].1, "What is Rust?");
964        assert_eq!(memories[0].2, "Rust is a systems language.");
965    }
966
967    // ========================================================================
968    // Integration-style Tests
969    // ========================================================================
970
971    #[tokio::test]
972    async fn test_multiple_providers_query() {
973        let provider1 = MockContextProvider::new("provider-1").with_items(vec![ContextItem::new(
974            "p1-1",
975            ContextType::Resource,
976            "Resource from P1",
977        )]);
978
979        let provider2 = MockContextProvider::new("provider-2").with_items(vec![
980            ContextItem::new("p2-1", ContextType::Memory, "Memory from P2"),
981            ContextItem::new("p2-2", ContextType::Skill, "Skill from P2"),
982        ]);
983
984        let providers: Vec<&dyn ContextProvider> = vec![&provider1, &provider2];
985        let query = ContextQuery::new("test");
986
987        let mut all_items = Vec::new();
988        for provider in providers {
989            let result = provider.query(&query).await.unwrap();
990            all_items.extend(result.items);
991        }
992
993        assert_eq!(all_items.len(), 3);
994        assert!(all_items.iter().any(|i| i.id == "p1-1"));
995        assert!(all_items.iter().any(|i| i.id == "p2-1"));
996        assert!(all_items.iter().any(|i| i.id == "p2-2"));
997    }
998
999    #[test]
1000    fn test_context_result_xml_formatting_complex() {
1001        let mut result = ContextResult::new("openviking");
1002        result.add_item(
1003            ContextItem::new(
1004                "doc-1",
1005                ContextType::Resource,
1006                "Authentication uses JWT tokens stored in httpOnly cookies.",
1007            )
1008            .with_source("viking://docs/auth")
1009            .with_token_count(50),
1010        );
1011        result.add_item(
1012            ContextItem::new(
1013                "mem-1",
1014                ContextType::Memory,
1015                "User prefers TypeScript over JavaScript.",
1016            )
1017            .with_token_count(30),
1018        );
1019
1020        let xml = result.to_xml();
1021
1022        // Verify structure
1023        assert!(xml.contains("<context source=\"viking://docs/auth\" type=\"Resource\">"));
1024        assert!(xml.contains("Authentication uses JWT tokens"));
1025        assert!(xml.contains("<context type=\"Memory\">"));
1026        assert!(xml.contains("User prefers TypeScript"));
1027
1028        // Verify items are separated
1029        assert!(xml.contains("</context>\n\n<context"));
1030    }
1031}
a3s_code_core/context/mod.rs

a3s_code_core/context/
mod.rs
