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

a3s_code_core/context/
mod.rs
