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 static_provider;
31
32pub use assembler::{
33    ContextAssembler, ContextAssembly, ContextAssemblyPolicy, ContextBudget, ContextSourcePolicy,
34};
35pub use fs_provider::{FileSystemContextConfig, FileSystemContextProvider};
36pub use recent_workspace_provider::RecentWorkspaceFilesContextProvider;
37pub use ripgrep_provider::{RipgrepContextConfig, RipgrepContextProvider};
38pub use static_provider::StaticContextProvider;
39
40use serde::{Deserialize, Serialize};
41use std::collections::HashMap;
42
43pub const CONTEXT_PROVENANCE_METADATA: &str = "a3s.context.provenance";
44pub const CONTEXT_PRIORITY_METADATA: &str = "a3s.context.priority";
45pub const CONTEXT_TRUST_METADATA: &str = "a3s.context.trust";
46pub const CONTEXT_FRESHNESS_METADATA: &str = "a3s.context.freshness";
47
48/// Type of context being queried
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
50pub enum ContextType {
51    /// Session/user history, extracted insights
52    Memory,
53    /// Documentation, code, knowledge base
54    #[default]
55    Resource,
56    /// Agent capabilities, behavior instructions
57    Skill,
58}
59
60/// Retrieval depth for tiered context (L0/L1/L2 pattern)
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
62pub enum ContextDepth {
63    /// ~100 tokens - high-level summary
64    Abstract,
65    /// ~2k tokens - key details (default)
66    #[default]
67    Overview,
68    /// Variable - complete content
69    Full,
70}
71
72/// Query to a context provider
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct ContextQuery {
75    /// The query string to search for relevant context
76    pub query: String,
77
78    /// Types of context to retrieve
79    #[serde(default)]
80    pub context_types: Vec<ContextType>,
81
82    /// Desired retrieval depth
83    #[serde(default)]
84    pub depth: ContextDepth,
85
86    /// Maximum number of results to return
87    #[serde(default = "default_max_results")]
88    pub max_results: usize,
89
90    /// Maximum total tokens across all results
91    #[serde(default = "default_max_tokens")]
92    pub max_tokens: usize,
93
94    /// Optional session ID for session-specific context
95    #[serde(default)]
96    pub session_id: Option<String>,
97
98    /// Additional provider-specific parameters
99    #[serde(default)]
100    pub params: HashMap<String, serde_json::Value>,
101}
102
103fn default_max_results() -> usize {
104    10
105}
106
107fn default_max_tokens() -> usize {
108    4000
109}
110
111impl ContextQuery {
112    /// Create a new context query with defaults
113    pub fn new(query: impl Into<String>) -> Self {
114        Self {
115            query: query.into(),
116            context_types: vec![ContextType::Resource],
117            depth: ContextDepth::default(),
118            max_results: default_max_results(),
119            max_tokens: default_max_tokens(),
120            session_id: None,
121            params: HashMap::new(),
122        }
123    }
124
125    /// Set the context types to retrieve
126    pub fn with_types(mut self, types: impl IntoIterator<Item = ContextType>) -> Self {
127        self.context_types = types.into_iter().collect();
128        self
129    }
130
131    /// Set the retrieval depth
132    pub fn with_depth(mut self, depth: ContextDepth) -> Self {
133        self.depth = depth;
134        self
135    }
136
137    /// Set the maximum number of results
138    pub fn with_max_results(mut self, max: usize) -> Self {
139        self.max_results = max;
140        self
141    }
142
143    /// Set the maximum total tokens
144    pub fn with_max_tokens(mut self, max: usize) -> Self {
145        self.max_tokens = max;
146        self
147    }
148
149    /// Set the session ID
150    pub fn with_session_id(mut self, id: impl Into<String>) -> Self {
151        self.session_id = Some(id.into());
152        self
153    }
154
155    /// Add a custom parameter
156    pub fn with_param(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
157        self.params.insert(key.into(), value);
158        self
159    }
160}
161
162/// A single piece of retrieved context
163#[derive(Debug, Clone, Serialize, Deserialize)]
164pub struct ContextItem {
165    /// Unique identifier for this context item
166    pub id: String,
167
168    /// Type of context
169    pub context_type: ContextType,
170
171    /// The actual content
172    pub content: String,
173
174    /// Estimated token count (informational)
175    #[serde(default)]
176    pub token_count: usize,
177
178    /// Relevance score (0.0 to 1.0)
179    #[serde(default)]
180    pub relevance: f32,
181
182    /// Optional source URI (e.g., "viking://docs/auth")
183    #[serde(default)]
184    pub source: Option<String>,
185
186    /// Additional metadata
187    #[serde(default)]
188    pub metadata: HashMap<String, serde_json::Value>,
189}
190
191impl ContextItem {
192    /// Create a new context item
193    pub fn new(
194        id: impl Into<String>,
195        context_type: ContextType,
196        content: impl Into<String>,
197    ) -> Self {
198        Self {
199            id: id.into(),
200            context_type,
201            content: content.into(),
202            token_count: 0,
203            relevance: 0.0,
204            source: None,
205            metadata: HashMap::new(),
206        }
207    }
208
209    /// Set the token count
210    pub fn with_token_count(mut self, count: usize) -> Self {
211        self.token_count = count;
212        self
213    }
214
215    /// Set the relevance score
216    pub fn with_relevance(mut self, score: f32) -> Self {
217        self.relevance = score.clamp(0.0, 1.0);
218        self
219    }
220
221    /// Set the source URI
222    pub fn with_source(mut self, source: impl Into<String>) -> Self {
223        self.source = Some(source.into());
224        self
225    }
226
227    /// Add metadata
228    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
229        self.metadata.insert(key.into(), value);
230        self
231    }
232
233    /// Set a human-readable provenance label for diagnostics and ranking.
234    pub fn with_provenance(mut self, provenance: impl Into<String>) -> Self {
235        self.metadata.insert(
236            CONTEXT_PROVENANCE_METADATA.to_string(),
237            serde_json::Value::String(provenance.into()),
238        );
239        self
240    }
241
242    /// Set priority score (0.0 to 1.0) for harness-controlled ranking.
243    pub fn with_priority(mut self, priority: f32) -> Self {
244        self.metadata.insert(
245            CONTEXT_PRIORITY_METADATA.to_string(),
246            serde_json::json!(priority.clamp(0.0, 1.0)),
247        );
248        self
249    }
250
251    /// Set trust score (0.0 to 1.0) for harness-controlled ranking.
252    pub fn with_trust(mut self, trust: f32) -> Self {
253        self.metadata.insert(
254            CONTEXT_TRUST_METADATA.to_string(),
255            serde_json::json!(trust.clamp(0.0, 1.0)),
256        );
257        self
258    }
259
260    /// Set freshness score (0.0 to 1.0) for harness-controlled ranking.
261    pub fn with_freshness(mut self, freshness: f32) -> Self {
262        self.metadata.insert(
263            CONTEXT_FRESHNESS_METADATA.to_string(),
264            serde_json::json!(freshness.clamp(0.0, 1.0)),
265        );
266        self
267    }
268
269    pub fn provenance(&self) -> Option<&str> {
270        self.metadata
271            .get(CONTEXT_PROVENANCE_METADATA)
272            .and_then(serde_json::Value::as_str)
273    }
274
275    pub fn priority(&self) -> f32 {
276        metadata_score(self.metadata.get(CONTEXT_PRIORITY_METADATA))
277    }
278
279    pub fn trust(&self) -> f32 {
280        metadata_score(self.metadata.get(CONTEXT_TRUST_METADATA))
281    }
282
283    pub fn freshness(&self) -> f32 {
284        metadata_score(self.metadata.get(CONTEXT_FRESHNESS_METADATA))
285    }
286
287    /// Format as XML tag for system prompt injection
288    pub fn to_xml(&self) -> String {
289        let source_attr = self
290            .source
291            .as_ref()
292            .map(|s| format!(" source=\"{}\"", s))
293            .unwrap_or_default();
294        let type_str = match self.context_type {
295            ContextType::Memory => "Memory",
296            ContextType::Resource => "Resource",
297            ContextType::Skill => "Skill",
298        };
299        format!(
300            "<context{} type=\"{}\">\n{}\n</context>",
301            source_attr, type_str, self.content
302        )
303    }
304}
305
306fn metadata_score(value: Option<&serde_json::Value>) -> f32 {
307    value
308        .and_then(serde_json::Value::as_f64)
309        .map(|score| (score as f32).clamp(0.0, 1.0))
310        .unwrap_or(0.0)
311}
312
313/// Result from a context provider query
314#[derive(Debug, Clone, Default, Serialize, Deserialize)]
315pub struct ContextResult {
316    /// Retrieved context items
317    pub items: Vec<ContextItem>,
318
319    /// Total tokens across all items
320    pub total_tokens: usize,
321
322    /// Name of the provider that returned these results
323    pub provider: String,
324
325    /// Whether results were truncated due to limits
326    pub truncated: bool,
327}
328
329impl ContextResult {
330    /// Create a new empty result
331    pub fn new(provider: impl Into<String>) -> Self {
332        Self {
333            items: Vec::new(),
334            total_tokens: 0,
335            provider: provider.into(),
336            truncated: false,
337        }
338    }
339
340    /// Add an item to the result
341    pub fn add_item(&mut self, item: ContextItem) {
342        self.total_tokens += item.token_count;
343        self.items.push(item);
344    }
345
346    /// Check if the result is empty
347    pub fn is_empty(&self) -> bool {
348        self.items.is_empty()
349    }
350
351    /// Format all items as XML for system prompt injection
352    pub fn to_xml(&self) -> String {
353        self.items
354            .iter()
355            .map(|item| item.to_xml())
356            .collect::<Vec<_>>()
357            .join("\n\n")
358    }
359}
360
361/// Context provider trait - implement this for OpenViking, RAG systems, etc.
362#[async_trait::async_trait]
363pub trait ContextProvider: Send + Sync {
364    /// Provider name (used for identification and logging)
365    fn name(&self) -> &str;
366
367    /// Query the provider for relevant context
368    async fn query(&self, query: &ContextQuery) -> anyhow::Result<ContextResult>;
369
370    /// Called after each turn for memory extraction (optional)
371    ///
372    /// Providers can implement this to extract and store memories from
373    /// the conversation.
374    async fn on_turn_complete(
375        &self,
376        _session_id: &str,
377        _prompt: &str,
378        _response: &str,
379    ) -> anyhow::Result<()> {
380        Ok(())
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    // ========================================================================
389    // ContextType Tests
390    // ========================================================================
391
392    #[test]
393    fn test_context_type_default() {
394        let ct: ContextType = Default::default();
395        assert_eq!(ct, ContextType::Resource);
396    }
397
398    #[test]
399    fn test_context_type_serialization() {
400        let ct = ContextType::Memory;
401        let json = serde_json::to_string(&ct).unwrap();
402        assert_eq!(json, "\"Memory\"");
403
404        let parsed: ContextType = serde_json::from_str(&json).unwrap();
405        assert_eq!(parsed, ContextType::Memory);
406    }
407
408    #[test]
409    fn test_context_type_all_variants() {
410        let types = vec![
411            ContextType::Memory,
412            ContextType::Resource,
413            ContextType::Skill,
414        ];
415        for ct in types {
416            let json = serde_json::to_string(&ct).unwrap();
417            let parsed: ContextType = serde_json::from_str(&json).unwrap();
418            assert_eq!(parsed, ct);
419        }
420    }
421
422    // ========================================================================
423    // ContextDepth Tests
424    // ========================================================================
425
426    #[test]
427    fn test_context_depth_default() {
428        let cd: ContextDepth = Default::default();
429        assert_eq!(cd, ContextDepth::Overview);
430    }
431
432    #[test]
433    fn test_context_depth_serialization() {
434        let cd = ContextDepth::Full;
435        let json = serde_json::to_string(&cd).unwrap();
436        assert_eq!(json, "\"Full\"");
437
438        let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
439        assert_eq!(parsed, ContextDepth::Full);
440    }
441
442    #[test]
443    fn test_context_depth_all_variants() {
444        let depths = vec![
445            ContextDepth::Abstract,
446            ContextDepth::Overview,
447            ContextDepth::Full,
448        ];
449        for cd in depths {
450            let json = serde_json::to_string(&cd).unwrap();
451            let parsed: ContextDepth = serde_json::from_str(&json).unwrap();
452            assert_eq!(parsed, cd);
453        }
454    }
455
456    // ========================================================================
457    // ContextQuery Tests
458    // ========================================================================
459
460    #[test]
461    fn test_context_query_new() {
462        let query = ContextQuery::new("test query");
463        assert_eq!(query.query, "test query");
464        assert_eq!(query.context_types, vec![ContextType::Resource]);
465        assert_eq!(query.depth, ContextDepth::Overview);
466        assert_eq!(query.max_results, 10);
467        assert_eq!(query.max_tokens, 4000);
468        assert!(query.session_id.is_none());
469        assert!(query.params.is_empty());
470    }
471
472    #[test]
473    fn test_context_query_builder() {
474        let query = ContextQuery::new("test")
475            .with_types([ContextType::Memory, ContextType::Skill])
476            .with_depth(ContextDepth::Full)
477            .with_max_results(5)
478            .with_max_tokens(2000)
479            .with_session_id("sess-123")
480            .with_param("custom", serde_json::json!("value"));
481
482        assert_eq!(query.context_types.len(), 2);
483        assert!(query.context_types.contains(&ContextType::Memory));
484        assert!(query.context_types.contains(&ContextType::Skill));
485        assert_eq!(query.depth, ContextDepth::Full);
486        assert_eq!(query.max_results, 5);
487        assert_eq!(query.max_tokens, 2000);
488        assert_eq!(query.session_id, Some("sess-123".to_string()));
489        assert_eq!(
490            query.params.get("custom"),
491            Some(&serde_json::json!("value"))
492        );
493    }
494
495    #[test]
496    fn test_context_query_serialization() {
497        let query = ContextQuery::new("search term")
498            .with_types([ContextType::Resource])
499            .with_session_id("sess-456");
500
501        let json = serde_json::to_string(&query).unwrap();
502        let parsed: ContextQuery = serde_json::from_str(&json).unwrap();
503
504        assert_eq!(parsed.query, "search term");
505        assert_eq!(parsed.session_id, Some("sess-456".to_string()));
506    }
507
508    #[test]
509    fn test_context_query_deserialization_with_defaults() {
510        let json = r#"{"query": "minimal query"}"#;
511        let query: ContextQuery = serde_json::from_str(json).unwrap();
512
513        assert_eq!(query.query, "minimal query");
514        assert!(query.context_types.is_empty()); // Default from serde is empty vec
515        assert_eq!(query.depth, ContextDepth::Overview);
516        assert_eq!(query.max_results, 10);
517        assert_eq!(query.max_tokens, 4000);
518    }
519
520    // ========================================================================
521    // ContextItem Tests
522    // ========================================================================
523
524    #[test]
525    fn test_context_item_new() {
526        let item = ContextItem::new("item-1", ContextType::Resource, "Some content");
527        assert_eq!(item.id, "item-1");
528        assert_eq!(item.context_type, ContextType::Resource);
529        assert_eq!(item.content, "Some content");
530        assert_eq!(item.token_count, 0);
531        assert_eq!(item.relevance, 0.0);
532        assert!(item.source.is_none());
533        assert!(item.metadata.is_empty());
534    }
535
536    #[test]
537    fn test_context_item_builder() {
538        let item = ContextItem::new("item-2", ContextType::Memory, "Memory content")
539            .with_token_count(150)
540            .with_relevance(0.85)
541            .with_source("viking://memory/session-123")
542            .with_provenance("memory")
543            .with_priority(0.7)
544            .with_trust(1.2)
545            .with_freshness(-1.0)
546            .with_metadata("key", serde_json::json!("value"));
547
548        assert_eq!(item.token_count, 150);
549        assert!((item.relevance - 0.85).abs() < f32::EPSILON);
550        assert_eq!(item.source, Some("viking://memory/session-123".to_string()));
551        assert_eq!(item.provenance(), Some("memory"));
552        assert!((item.priority() - 0.7).abs() < f32::EPSILON);
553        assert!((item.trust() - 1.0).abs() < f32::EPSILON);
554        assert!(item.freshness().abs() < f32::EPSILON);
555        assert_eq!(item.metadata.get("key"), Some(&serde_json::json!("value")));
556    }
557
558    #[test]
559    fn test_context_item_relevance_clamping() {
560        let item1 = ContextItem::new("id", ContextType::Resource, "").with_relevance(1.5);
561        assert!((item1.relevance - 1.0).abs() < f32::EPSILON);
562
563        let item2 = ContextItem::new("id", ContextType::Resource, "").with_relevance(-0.5);
564        assert!(item2.relevance.abs() < f32::EPSILON);
565    }
566
567    #[test]
568    fn test_context_item_to_xml_without_source() {
569        let item = ContextItem::new("id", ContextType::Resource, "Content here");
570        let xml = item.to_xml();
571        assert_eq!(xml, "<context type=\"Resource\">\nContent here\n</context>");
572    }
573
574    #[test]
575    fn test_context_item_to_xml_with_source() {
576        let item = ContextItem::new("id", ContextType::Memory, "Memory content")
577            .with_source("viking://docs/auth");
578        let xml = item.to_xml();
579        assert_eq!(
580            xml,
581            "<context source=\"viking://docs/auth\" type=\"Memory\">\nMemory content\n</context>"
582        );
583    }
584
585    #[test]
586    fn test_context_item_to_xml_all_types() {
587        let memory = ContextItem::new("m", ContextType::Memory, "m").to_xml();
588        assert!(memory.contains("type=\"Memory\""));
589
590        let resource = ContextItem::new("r", ContextType::Resource, "r").to_xml();
591        assert!(resource.contains("type=\"Resource\""));
592
593        let skill = ContextItem::new("s", ContextType::Skill, "s").to_xml();
594        assert!(skill.contains("type=\"Skill\""));
595    }
596
597    #[test]
598    fn test_context_item_serialization() {
599        let item = ContextItem::new("item-3", ContextType::Skill, "Skill instructions")
600            .with_token_count(200)
601            .with_relevance(0.9)
602            .with_source("viking://skills/code-review");
603
604        let json = serde_json::to_string(&item).unwrap();
605        let parsed: ContextItem = serde_json::from_str(&json).unwrap();
606
607        assert_eq!(parsed.id, "item-3");
608        assert_eq!(parsed.context_type, ContextType::Skill);
609        assert_eq!(parsed.content, "Skill instructions");
610        assert_eq!(parsed.token_count, 200);
611    }
612
613    // ========================================================================
614    // ContextResult Tests
615    // ========================================================================
616
617    #[test]
618    fn test_context_result_new() {
619        let result = ContextResult::new("test-provider");
620        assert!(result.items.is_empty());
621        assert_eq!(result.total_tokens, 0);
622        assert_eq!(result.provider, "test-provider");
623        assert!(!result.truncated);
624    }
625
626    #[test]
627    fn test_context_result_add_item() {
628        let mut result = ContextResult::new("provider");
629        let item = ContextItem::new("id", ContextType::Resource, "content").with_token_count(100);
630        result.add_item(item);
631
632        assert_eq!(result.items.len(), 1);
633        assert_eq!(result.total_tokens, 100);
634    }
635
636    #[test]
637    fn test_context_result_add_multiple_items() {
638        let mut result = ContextResult::new("provider");
639        result.add_item(ContextItem::new("1", ContextType::Resource, "a").with_token_count(50));
640        result.add_item(ContextItem::new("2", ContextType::Memory, "b").with_token_count(75));
641        result.add_item(ContextItem::new("3", ContextType::Skill, "c").with_token_count(25));
642
643        assert_eq!(result.items.len(), 3);
644        assert_eq!(result.total_tokens, 150);
645    }
646
647    #[test]
648    fn test_context_result_is_empty() {
649        let empty = ContextResult::new("provider");
650        assert!(empty.is_empty());
651
652        let mut non_empty = ContextResult::new("provider");
653        non_empty.add_item(ContextItem::new("id", ContextType::Resource, "content"));
654        assert!(!non_empty.is_empty());
655    }
656
657    #[test]
658    fn test_context_result_to_xml() {
659        let mut result = ContextResult::new("provider");
660        result.add_item(
661            ContextItem::new("1", ContextType::Resource, "First content").with_source("source://1"),
662        );
663        result.add_item(ContextItem::new("2", ContextType::Memory, "Second content"));
664
665        let xml = result.to_xml();
666        assert!(xml.contains("<context source=\"source://1\" type=\"Resource\">"));
667        assert!(xml.contains("First content"));
668        assert!(xml.contains("<context type=\"Memory\">"));
669        assert!(xml.contains("Second content"));
670    }
671
672    #[test]
673    fn test_context_result_to_xml_empty() {
674        let result = ContextResult::new("provider");
675        let xml = result.to_xml();
676        assert!(xml.is_empty());
677    }
678
679    #[test]
680    fn test_context_result_serialization() {
681        let mut result = ContextResult::new("test-provider");
682        result.truncated = true;
683        result.add_item(ContextItem::new("id", ContextType::Resource, "content"));
684
685        let json = serde_json::to_string(&result).unwrap();
686        let parsed: ContextResult = serde_json::from_str(&json).unwrap();
687
688        assert_eq!(parsed.provider, "test-provider");
689        assert!(parsed.truncated);
690        assert_eq!(parsed.items.len(), 1);
691    }
692
693    #[test]
694    fn test_context_result_default() {
695        let result: ContextResult = Default::default();
696        assert!(result.items.is_empty());
697        assert_eq!(result.total_tokens, 0);
698        assert!(result.provider.is_empty());
699        assert!(!result.truncated);
700    }
701
702    // ========================================================================
703    // ContextProvider Trait Tests (with Mock)
704    // ========================================================================
705
706    struct MockContextProvider {
707        name: String,
708        items: Vec<ContextItem>,
709    }
710
711    impl MockContextProvider {
712        fn new(name: &str) -> Self {
713            Self {
714                name: name.to_string(),
715                items: Vec::new(),
716            }
717        }
718
719        fn with_items(mut self, items: Vec<ContextItem>) -> Self {
720            self.items = items;
721            self
722        }
723    }
724
725    #[async_trait::async_trait]
726    impl ContextProvider for MockContextProvider {
727        fn name(&self) -> &str {
728            &self.name
729        }
730
731        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
732            let mut result = ContextResult::new(&self.name);
733            for item in &self.items {
734                result.add_item(item.clone());
735            }
736            Ok(result)
737        }
738    }
739
740    #[tokio::test]
741    async fn test_mock_context_provider() {
742        let provider = MockContextProvider::new("mock").with_items(vec![ContextItem::new(
743            "1",
744            ContextType::Resource,
745            "content",
746        )]);
747
748        assert_eq!(provider.name(), "mock");
749
750        let query = ContextQuery::new("test");
751        let result = provider.query(&query).await.unwrap();
752
753        assert_eq!(result.provider, "mock");
754        assert_eq!(result.items.len(), 1);
755    }
756
757    #[tokio::test]
758    async fn test_context_provider_on_turn_complete_default() {
759        let provider = MockContextProvider::new("mock");
760
761        // Default implementation should succeed
762        let result = provider
763            .on_turn_complete("session-1", "prompt", "response")
764            .await;
765        assert!(result.is_ok());
766    }
767
768    struct MockMemoryProvider {
769        memories: std::sync::Arc<tokio::sync::RwLock<Vec<(String, String, String)>>>,
770    }
771
772    impl MockMemoryProvider {
773        fn new() -> Self {
774            Self {
775                memories: std::sync::Arc::new(tokio::sync::RwLock::new(Vec::new())),
776            }
777        }
778    }
779
780    #[async_trait::async_trait]
781    impl ContextProvider for MockMemoryProvider {
782        fn name(&self) -> &str {
783            "memory-provider"
784        }
785
786        async fn query(&self, _query: &ContextQuery) -> anyhow::Result<ContextResult> {
787            Ok(ContextResult::new("memory-provider"))
788        }
789
790        async fn on_turn_complete(
791            &self,
792            session_id: &str,
793            prompt: &str,
794            response: &str,
795        ) -> anyhow::Result<()> {
796            let mut memories = self.memories.write().await;
797            memories.push((
798                session_id.to_string(),
799                prompt.to_string(),
800                response.to_string(),
801            ));
802            Ok(())
803        }
804    }
805
806    #[tokio::test]
807    async fn test_context_provider_on_turn_complete_custom() {
808        let provider = MockMemoryProvider::new();
809
810        provider
811            .on_turn_complete("sess-1", "What is Rust?", "Rust is a systems language.")
812            .await
813            .unwrap();
814
815        let memories = provider.memories.read().await;
816        assert_eq!(memories.len(), 1);
817        assert_eq!(memories[0].0, "sess-1");
818        assert_eq!(memories[0].1, "What is Rust?");
819        assert_eq!(memories[0].2, "Rust is a systems language.");
820    }
821
822    // ========================================================================
823    // Integration-style Tests
824    // ========================================================================
825
826    #[tokio::test]
827    async fn test_multiple_providers_query() {
828        let provider1 = MockContextProvider::new("provider-1").with_items(vec![ContextItem::new(
829            "p1-1",
830            ContextType::Resource,
831            "Resource from P1",
832        )]);
833
834        let provider2 = MockContextProvider::new("provider-2").with_items(vec![
835            ContextItem::new("p2-1", ContextType::Memory, "Memory from P2"),
836            ContextItem::new("p2-2", ContextType::Skill, "Skill from P2"),
837        ]);
838
839        let providers: Vec<&dyn ContextProvider> = vec![&provider1, &provider2];
840        let query = ContextQuery::new("test");
841
842        let mut all_items = Vec::new();
843        for provider in providers {
844            let result = provider.query(&query).await.unwrap();
845            all_items.extend(result.items);
846        }
847
848        assert_eq!(all_items.len(), 3);
849        assert!(all_items.iter().any(|i| i.id == "p1-1"));
850        assert!(all_items.iter().any(|i| i.id == "p2-1"));
851        assert!(all_items.iter().any(|i| i.id == "p2-2"));
852    }
853
854    #[test]
855    fn test_context_result_xml_formatting_complex() {
856        let mut result = ContextResult::new("openviking");
857        result.add_item(
858            ContextItem::new(
859                "doc-1",
860                ContextType::Resource,
861                "Authentication uses JWT tokens stored in httpOnly cookies.",
862            )
863            .with_source("viking://docs/auth")
864            .with_token_count(50),
865        );
866        result.add_item(
867            ContextItem::new(
868                "mem-1",
869                ContextType::Memory,
870                "User prefers TypeScript over JavaScript.",
871            )
872            .with_token_count(30),
873        );
874
875        let xml = result.to_xml();
876
877        // Verify structure
878        assert!(xml.contains("<context source=\"viking://docs/auth\" type=\"Resource\">"));
879        assert!(xml.contains("Authentication uses JWT tokens"));
880        assert!(xml.contains("<context type=\"Memory\">"));
881        assert!(xml.contains("User prefers TypeScript"));
882
883        // Verify items are separated
884        assert!(xml.contains("</context>\n\n<context"));
885    }
886}
a3s_code_core/context/mod.rs

a3s_code_core/context/
mod.rs
