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=\"{}\"", 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, type_str, self.content
325        )
326    }
327}
328
329fn metadata_score(value: Option<&serde_json::Value>) -> f32 {
330    value
331        .and_then(serde_json::Value::as_f64)
332        .map(|score| (score as f32).clamp(0.0, 1.0))
333        .unwrap_or(0.0)
334}
335
336/// Result from a context provider query
337#[derive(Debug, Clone, Default, Serialize, Deserialize)]
338pub struct ContextResult {
339    /// Retrieved context items
340    pub items: Vec<ContextItem>,
341
342    /// Total tokens across all items
343    pub total_tokens: usize,
344
345    /// Name of the provider that returned these results
346    pub provider: String,
347
348    /// Whether results were truncated due to limits
349    pub truncated: bool,
350}
351
352/// Runtime behavior when one context provider returns an error.
353#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
354pub enum ContextProviderFailureMode {
355    /// Preserve the existing general-purpose RAG behavior: log and omit the
356    /// failed provider while other sources may continue.
357    #[default]
358    BestEffort,
359    /// Abort the turn. Used by exact-generation cognitive packages so a
360    /// provider failure cannot silently fall back to memory, graph data, or an
361    /// unpinned package.
362    FailClosed,
363}
364
365impl ContextResult {
366    /// Create a new empty result
367    pub fn new(provider: impl Into<String>) -> Self {
368        Self {
369            items: Vec::new(),
370            total_tokens: 0,
371            provider: provider.into(),
372            truncated: false,
373        }
374    }
375
376    /// Add an item to the result
377    pub fn add_item(&mut self, item: ContextItem) {
378        self.total_tokens += item.token_count;
379        self.items.push(item);
380    }
381
382    /// Check if the result is empty
383    pub fn is_empty(&self) -> bool {
384        self.items.is_empty()
385    }
386
387    /// Format all items as XML for system prompt injection
388    pub fn to_xml(&self) -> String {
389        self.items
390            .iter()
391            .map(|item| item.to_xml())
392            .collect::<Vec<_>>()
393            .join("\n\n")
394    }
395}
396
397/// Context provider trait - implement this for OpenViking, RAG systems, etc.
398#[async_trait::async_trait]
399pub trait ContextProvider: Send + Sync {
400    /// Provider name (used for identification and logging)
401    fn name(&self) -> &str;
402
403    /// Whether provider failure may be omitted or must abort the turn.
404    fn failure_mode(&self) -> ContextProviderFailureMode {
405        ContextProviderFailureMode::BestEffort
406    }
407
408    /// Exact cognitive-package identity, when this is Code's typed adapter.
409    ///
410    /// Hosts must install such an adapter through
411    /// [`SessionOptions::with_cognitive_context`](crate::SessionOptions::with_cognitive_context)
412    /// so the same value is persisted in the session snapshot.
413    fn cognitive_package_binding(
414        &self,
415    ) -> Option<&crate::cognitive_context::CognitivePackageBindingV1> {
416        None
417    }
418
419    /// Query the provider for relevant context
420    async fn query(&self, query: &ContextQuery) -> anyhow::Result<ContextResult>;
421
422    /// Called after each turn for memory extraction (optional)
423    ///
424    /// Providers can implement this to extract and store memories from
425    /// the conversation.
426    async fn on_turn_complete(
427        &self,
428        _session_id: &str,
429        _prompt: &str,
430        _response: &str,
431    ) -> anyhow::Result<()> {
432        Ok(())
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    // ========================================================================
441    // ContextType Tests
442    // ========================================================================
443
444    #[test]
445    fn test_context_type_default() {
446        let ct: ContextType = Default::default();
447        assert_eq!(ct, ContextType::Resource);
448    }
449
450    #[test]
451    fn test_context_type_serialization() {
452        let ct = ContextType::Memory;
453        let json = serde_json::to_string(&ct).unwrap();
454        assert_eq!(json, "\"Memory\"");
455
456        let parsed: ContextType = serde_json::from_str(&json).unwrap();
457        assert_eq!(parsed, ContextType::Memory);
458    }
459
460    #[test]
461    fn test_context_type_all_variants() {
462        let types = vec![
463            ContextType::Memory,
464            ContextType::Resource,
465            ContextType::Skill,
466        ];
467        for ct in types {
468            let json = serde_json::to_string(&ct).unwrap();
469            let parsed: ContextType = serde_json::from_str(&json).unwrap();
470            assert_eq!(parsed, ct);
471        }
472    }
473
474    // ========================================================================
475    // ContextDepth Tests
476    // ========================================================================
477
478    #[test]
479    fn test_context_depth_default() {
480        let cd: ContextDepth = Default::default();
481        assert_eq!(cd, ContextDepth::Overview);
482    }
483
484    #[test]
485    fn test_context_depth_serialization() {
486        let cd = ContextDepth::Full;
487        let json = serde_json::to_string(&cd).unwrap();
488        assert_eq!(json, "\"Full\"");
489
490        let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
491        assert_eq!(parsed, ContextDepth::Full);
492    }
493
494    #[test]
495    fn test_context_depth_all_variants() {
496        let depths = vec![
497            ContextDepth::Abstract,
498            ContextDepth::Overview,
499            ContextDepth::Full,
500        ];
501        for cd in depths {
502            let json = serde_json::to_string(&cd).unwrap();
503            let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
504            assert_eq!(parsed, cd);
505        }
506    }
507
508    // ========================================================================
509    // ContextQuery Tests
510    // ========================================================================
511
512    #[test]
513    fn test_context_query_new() {
514        let query = ContextQuery::new("test query");
515        assert_eq!(query.query, "test query");
516        assert_eq!(query.context_types, vec![ContextType::Resource]);
517        assert_eq!(query.depth, ContextDepth::Overview);
518        assert_eq!(query.max_results, 10);
519        assert_eq!(query.max_tokens, 4000);
520        assert!(query.session_id.is_none());
521        assert!(query.params.is_empty());
522    }
523
524    #[test]
525    fn test_context_query_builder() {
526        let query = ContextQuery::new("test")
527            .with_types([ContextType::Memory, ContextType::Skill])
528            .with_depth(ContextDepth::Full)
529            .with_max_results(5)
530            .with_max_tokens(2000)
531            .with_session_id("sess-123")
532            .with_param("custom", serde_json::json!("value"));
533
534        assert_eq!(query.context_types.len(), 2);
535        assert!(query.context_types.contains(&ContextType::Memory));
536        assert!(query.context_types.contains(&ContextType::Skill));
537        assert_eq!(query.depth, ContextDepth::Full);
538        assert_eq!(query.max_results, 5);
539        assert_eq!(query.max_tokens, 2000);
540        assert_eq!(query.session_id, Some("sess-123".to_string()));
541        assert_eq!(
542            query.params.get("custom"),
543            Some(&serde_json::json!("value"))
544        );
545    }
546
547    #[test]
548    fn test_context_query_serialization() {
549        let query = ContextQuery::new("search term")
550            .with_types([ContextType::Resource])
551            .with_session_id("sess-456");
552
553        let json = serde_json::to_string(&query).unwrap();
554        let parsed: ContextQuery = serde_json::from_str(&json).unwrap();
555
556        assert_eq!(parsed.query, "search term");
557        assert_eq!(parsed.session_id, Some("sess-456".to_string()));
558    }
559
560    #[test]
561    fn test_context_query_deserialization_with_defaults() {
562        let json = r#"{"query": "minimal query"}"#;
563        let query: ContextQuery = serde_json::from_str(json).unwrap();
564
565        assert_eq!(query.query, "minimal query");
566        assert!(query.context_types.is_empty()); // Default from serde is empty vec
567        assert_eq!(query.depth, ContextDepth::Overview);
568        assert_eq!(query.max_results, 10);
569        assert_eq!(query.max_tokens, 4000);
570    }
571
572    // ========================================================================
573    // ContextItem Tests
574    // ========================================================================
575
576    #[test]
577    fn test_context_item_new() {
578        let item = ContextItem::new("item-1", ContextType::Resource, "Some content");
579        assert_eq!(item.id, "item-1");
580        assert_eq!(item.context_type, ContextType::Resource);
581        assert_eq!(item.content, "Some content");
582        assert_eq!(item.token_count, 0);
583        assert_eq!(item.relevance, 0.0);
584        assert!(item.source.is_none());
585        assert!(item.metadata.is_empty());
586    }
587
588    #[test]
589    fn test_context_item_builder() {
590        let item = ContextItem::new("item-2", ContextType::Memory, "Memory content")
591            .with_token_count(150)
592            .with_relevance(0.85)
593            .with_source("viking://memory/session-123")
594            .with_provenance("memory")
595            .with_priority(0.7)
596            .with_trust(1.2)
597            .with_freshness(-1.0)
598            .with_metadata("key", serde_json::json!("value"));
599
600        assert_eq!(item.token_count, 150);
601        assert!((item.relevance - 0.85).abs() < f32::EPSILON);
602        assert_eq!(item.source, Some("viking://memory/session-123".to_string()));
603        assert_eq!(item.provenance(), Some("memory"));
604        assert!((item.priority() - 0.7).abs() < f32::EPSILON);
605        assert!((item.trust() - 1.0).abs() < f32::EPSILON);
606        assert!(item.freshness().abs() < f32::EPSILON);
607        assert_eq!(item.metadata.get("key"), Some(&serde_json::json!("value")));
608    }
609
610    #[test]
611    fn test_context_item_relevance_clamping() {
612        let item1 = ContextItem::new("id", ContextType::Resource, "").with_relevance(1.5);
613        assert!((item1.relevance - 1.0).abs() < f32::EPSILON);
614
615        let item2 = ContextItem::new("id", ContextType::Resource, "").with_relevance(-0.5);
616        assert!(item2.relevance.abs() < f32::EPSILON);
617    }
618
619    #[test]
620    fn test_context_item_to_xml_without_source() {
621        let item = ContextItem::new("id", ContextType::Resource, "Content here");
622        let xml = item.to_xml();
623        assert_eq!(xml, "<context type=\"Resource\">\nContent here\n</context>");
624    }
625
626    #[test]
627    fn test_context_item_to_xml_with_source() {
628        let item = ContextItem::new("id", ContextType::Memory, "Memory content")
629            .with_source("viking://docs/auth");
630        let xml = item.to_xml();
631        assert_eq!(
632            xml,
633            "<context source=\"viking://docs/auth\" type=\"Memory\">\nMemory content\n</context>"
634        );
635    }
636
637    #[test]
638    fn test_context_item_to_xml_all_types() {
639        let memory = ContextItem::new("m", ContextType::Memory, "m").to_xml();
640        assert!(memory.contains("type=\"Memory\""));
641
642        let resource = ContextItem::new("r", ContextType::Resource, "r").to_xml();
643        assert!(resource.contains("type=\"Resource\""));
644
645        let skill = ContextItem::new("s", ContextType::Skill, "s").to_xml();
646        assert!(skill.contains("type=\"Skill\""));
647    }
648
649    #[test]
650    fn test_context_item_serialization() {
651        let item = ContextItem::new("item-3", ContextType::Skill, "Skill instructions")
652            .with_token_count(200)
653            .with_relevance(0.9)
654            .with_source("viking://skills/code-review");
655
656        let json = serde_json::to_string(&item).unwrap();
657        let parsed: ContextItem = serde_json::from_str(&json).unwrap();
658
659        assert_eq!(parsed.id, "item-3");
660        assert_eq!(parsed.context_type, ContextType::Skill);
661        assert_eq!(parsed.content, "Skill instructions");
662        assert_eq!(parsed.token_count, 200);
663    }
664
665    // ========================================================================
666    // ContextResult Tests
667    // ========================================================================
668
669    #[test]
670    fn test_context_result_new() {
671        let result = ContextResult::new("test-provider");
672        assert!(result.items.is_empty());
673        assert_eq!(result.total_tokens, 0);
674        assert_eq!(result.provider, "test-provider");
675        assert!(!result.truncated);
676    }
677
678    #[test]
679    fn test_context_result_add_item() {
680        let mut result = ContextResult::new("provider");
681        let item = ContextItem::new("id", ContextType::Resource, "content").with_token_count(100);
682        result.add_item(item);
683
684        assert_eq!(result.items.len(), 1);
685        assert_eq!(result.total_tokens, 100);
686    }
687
688    #[test]
689    fn test_context_result_add_multiple_items() {
690        let mut result = ContextResult::new("provider");
691        result.add_item(ContextItem::new("1", ContextType::Resource, "a").with_token_count(50));
692        result.add_item(ContextItem::new("2", ContextType::Memory, "b").with_token_count(75));
693        result.add_item(ContextItem::new("3", ContextType::Skill, "c").with_token_count(25));
694
695        assert_eq!(result.items.len(), 3);
696        assert_eq!(result.total_tokens, 150);
697    }
698
699    #[test]
700    fn test_context_result_is_empty() {
701        let empty = ContextResult::new("provider");
702        assert!(empty.is_empty());
703
704        let mut non_empty = ContextResult::new("provider");
705        non_empty.add_item(ContextItem::new("id", ContextType::Resource, "content"));
706        assert!(!non_empty.is_empty());
707    }
708
709    #[test]
710    fn test_context_result_to_xml() {
711        let mut result = ContextResult::new("provider");
712        result.add_item(
713            ContextItem::new("1", ContextType::Resource, "First content").with_source("source://1"),
714        );
715        result.add_item(ContextItem::new("2", ContextType::Memory, "Second content"));
716
717        let xml = result.to_xml();
718        assert!(xml.contains("<context source=\"source://1\" type=\"Resource\">"));
719        assert!(xml.contains("First content"));
720        assert!(xml.contains("<context type=\"Memory\">"));
721        assert!(xml.contains("Second content"));
722    }
723
724    #[test]
725    fn test_context_result_to_xml_empty() {
726        let result = ContextResult::new("provider");
727        let xml = result.to_xml();
728        assert!(xml.is_empty());
729    }
730
731    #[test]
732    fn test_context_result_serialization() {
733        let mut result = ContextResult::new("test-provider");
734        result.truncated = true;
735        result.add_item(ContextItem::new("id", ContextType::Resource, "content"));
736
737        let json = serde_json::to_string(&result).unwrap();
738        let parsed: ContextResult = serde_json::from_str(&json).unwrap();
739
740        assert_eq!(parsed.provider, "test-provider");
741        assert!(parsed.truncated);
742        assert_eq!(parsed.items.len(), 1);
743    }
744
745    #[test]
746    fn test_context_result_default() {
747        let result: ContextResult = Default::default();
748        assert!(result.items.is_empty());
749        assert_eq!(result.total_tokens, 0);
750        assert!(result.provider.is_empty());
751        assert!(!result.truncated);
752    }
753
754    // ========================================================================
755    // ContextProvider Trait Tests (with Mock)
756    // ========================================================================
757
758    struct MockContextProvider {
759        name: String,
760        items: Vec<ContextItem>,
761    }
762
763    impl MockContextProvider {
764        fn new(name: &str) -> Self {
765            Self {
766                name: name.to_string(),
767                items: Vec::new(),
768            }
769        }
770
771        fn with_items(mut self, items: Vec<ContextItem>) -> Self {
772            self.items = items;
773            self
774        }
775    }
776
777    #[async_trait::async_trait]
778    impl ContextProvider for MockContextProvider {
779        fn name(&self) -> &str {
780            &self.name
781        }
782
783        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
784            let mut result = ContextResult::new(&self.name);
785            for item in &self.items {
786                result.add_item(item.clone());
787            }
788            Ok(result)
789        }
790    }
791
792    #[tokio::test]
793    async fn test_mock_context_provider() {
794        let provider = MockContextProvider::new("mock").with_items(vec![ContextItem::new(
795            "1",
796            ContextType::Resource,
797            "content",
798        )]);
799
800        assert_eq!(provider.name(), "mock");
801
802        let query = ContextQuery::new("test");
803        let result = provider.query(&query).await.unwrap();
804
805        assert_eq!(result.provider, "mock");
806        assert_eq!(result.items.len(), 1);
807    }
808
809    #[tokio::test]
810    async fn test_context_provider_on_turn_complete_default() {
811        let provider = MockContextProvider::new("mock");
812
813        // Default implementation should succeed
814        let result = provider
815            .on_turn_complete("session-1", "prompt", "response")
816            .await;
817        assert!(result.is_ok());
818    }
819
820    struct MockMemoryProvider {
821        memories: std::sync::Arc<tokio::sync::RwLock<Vec<(String, String, String)>>>,
822    }
823
824    impl MockMemoryProvider {
825        fn new() -> Self {
826            Self {
827                memories: std::sync::Arc::new(tokio::sync::RwLock::new(Vec::new())),
828            }
829        }
830    }
831
832    #[async_trait::async_trait]
833    impl ContextProvider for MockMemoryProvider {
834        fn name(&self) -> &str {
835            "memory-provider"
836        }
837
838        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
839            Ok(ContextResult::new("memory-provider"))
840        }
841
842        async fn on_turn_complete(
843            &self,
844            session_id: &str,
845            prompt: &str,
846            response: &str,
847        ) -> anyhow::Result<()> {
848            let mut memories = self.memories.write().await;
849            memories.push((
850                session_id.to_string(),
851                prompt.to_string(),
852                response.to_string(),
853            ));
854            Ok(())
855        }
856    }
857
858    #[tokio::test]
859    async fn test_context_provider_on_turn_complete_custom() {
860        let provider = MockMemoryProvider::new();
861
862        provider
863            .on_turn_complete("sess-1", "What is Rust?", "Rust is a systems language.")
864            .await
865            .unwrap();
866
867        let memories = provider.memories.read().await;
868        assert_eq!(memories.len(), 1);
869        assert_eq!(memories[0].0, "sess-1");
870        assert_eq!(memories[0].1, "What is Rust?");
871        assert_eq!(memories[0].2, "Rust is a systems language.");
872    }
873
874    // ========================================================================
875    // Integration-style Tests
876    // ========================================================================
877
878    #[tokio::test]
879    async fn test_multiple_providers_query() {
880        let provider1 = MockContextProvider::new("provider-1").with_items(vec![ContextItem::new(
881            "p1-1",
882            ContextType::Resource,
883            "Resource from P1",
884        )]);
885
886        let provider2 = MockContextProvider::new("provider-2").with_items(vec![
887            ContextItem::new("p2-1", ContextType::Memory, "Memory from P2"),
888            ContextItem::new("p2-2", ContextType::Skill, "Skill from P2"),
889        ]);
890
891        let providers: Vec<&dyn ContextProvider> = vec![&provider1, &provider2];
892        let query = ContextQuery::new("test");
893
894        let mut all_items = Vec::new();
895        for provider in providers {
896            let result = provider.query(&query).await.unwrap();
897            all_items.extend(result.items);
898        }
899
900        assert_eq!(all_items.len(), 3);
901        assert!(all_items.iter().any(|i| i.id == "p1-1"));
902        assert!(all_items.iter().any(|i| i.id == "p2-1"));
903        assert!(all_items.iter().any(|i| i.id == "p2-2"));
904    }
905
906    #[test]
907    fn test_context_result_xml_formatting_complex() {
908        let mut result = ContextResult::new("openviking");
909        result.add_item(
910            ContextItem::new(
911                "doc-1",
912                ContextType::Resource,
913                "Authentication uses JWT tokens stored in httpOnly cookies.",
914            )
915            .with_source("viking://docs/auth")
916            .with_token_count(50),
917        );
918        result.add_item(
919            ContextItem::new(
920                "mem-1",
921                ContextType::Memory,
922                "User prefers TypeScript over JavaScript.",
923            )
924            .with_token_count(30),
925        );
926
927        let xml = result.to_xml();
928
929        // Verify structure
930        assert!(xml.contains("<context source=\"viking://docs/auth\" type=\"Resource\">"));
931        assert!(xml.contains("Authentication uses JWT tokens"));
932        assert!(xml.contains("<context type=\"Memory\">"));
933        assert!(xml.contains("User prefers TypeScript"));
934
935        // Verify items are separated
936        assert!(xml.contains("</context>\n\n<context"));
937    }
938}
a3s_code_core/context/mod.rs

a3s_code_core/context/
mod.rs
