Skip to main content

relay_knowledge/domain/code/
context.rs

1use serde::{Deserialize, Serialize};
2
3use super::{
4    CodeQueryKind, CodeRepositorySelector, CodeRetrievalHit, CodeRetrievalLayer, DomainError,
5    FreshnessPolicy, RepositoryCodeRange, error::required_text,
6};
7
8pub const CODEGRAPH_CONTEXT_DEFAULT_LIMIT: usize = 8;
9pub const CODEGRAPH_CONTEXT_MAX_LIMIT: usize = 20;
10pub const CODEGRAPH_CONTEXT_MIN_BYTES: usize = 1024;
11pub const CODEGRAPH_CONTEXT_DEFAULT_MAX_BYTES: usize = 65_536;
12pub const CODEGRAPH_CONTEXT_MAX_BYTES: usize = 262_144;
13
14/// Agent-facing one-call code graph context request.
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct CodeGraphContextRequest {
17    pub repository: CodeRepositorySelector,
18    pub query: String,
19    pub limit: usize,
20    pub freshness_policy: FreshnessPolicy,
21    pub max_context_bytes: usize,
22    #[serde(default = "default_include_code")]
23    pub include_code: bool,
24    #[serde(default)]
25    pub exclude_generated: bool,
26}
27
28impl CodeGraphContextRequest {
29    /// Validates text and hard bounds before context orchestration starts.
30    pub fn new(
31        repository: CodeRepositorySelector,
32        query: impl Into<String>,
33        limit: usize,
34        freshness_policy: FreshnessPolicy,
35        max_context_bytes: usize,
36        include_code: bool,
37        exclude_generated: bool,
38    ) -> Result<Self, DomainError> {
39        let limit = match limit {
40            1..=CODEGRAPH_CONTEXT_MAX_LIMIT => limit,
41            0 => return Err(DomainError::invalid("limit", "must be greater than zero")),
42            _ => {
43                return Err(DomainError::invalid(
44                    "limit",
45                    "must be 20 or less for codegraph context",
46                ));
47            }
48        };
49        let max_context_bytes = match max_context_bytes {
50            0..CODEGRAPH_CONTEXT_MIN_BYTES => {
51                return Err(DomainError::invalid(
52                    "max_context_bytes",
53                    "must be at least 1024 for codegraph context",
54                ));
55            }
56            CODEGRAPH_CONTEXT_MIN_BYTES..=CODEGRAPH_CONTEXT_MAX_BYTES => max_context_bytes,
57            _ => {
58                return Err(DomainError::invalid(
59                    "max_context_bytes",
60                    "must be 262144 or less",
61                ));
62            }
63        };
64
65        Ok(Self {
66            repository,
67            query: required_text("query", query)?,
68            limit,
69            freshness_policy,
70            max_context_bytes,
71            include_code,
72            exclude_generated,
73        })
74    }
75}
76
77fn default_include_code() -> bool {
78    true
79}
80
81/// Context item provenance used by packed codegraph responses.
82#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
83pub struct CodeGraphContextProvenance {
84    pub query_kind: CodeQueryKind,
85    pub retrieval_layers: Vec<CodeRetrievalLayer>,
86    pub score: f64,
87}
88
89/// Compact code excerpt retained inside the context byte budget.
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
91pub struct CodeGraphCodeExcerpt {
92    pub path: String,
93    pub language_id: String,
94    pub line_range: RepositoryCodeRange,
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub symbol_snapshot_id: Option<String>,
97    pub provenance: CodeGraphContextProvenance,
98    pub excerpt: String,
99}
100
101/// Structural risk hint around a seed symbol or file.
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct CodeGraphImpactHint {
104    pub path: String,
105    pub line_range: RepositoryCodeRange,
106    pub relationship: String,
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub symbol_snapshot_id: Option<String>,
109    pub retrieval_layers: Vec<CodeRetrievalLayer>,
110    pub score: f64,
111}
112
113/// Budget consumed by one codegraph context orchestration.
114#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
115pub struct CodeGraphContextBudget {
116    pub limit: usize,
117    pub max_context_bytes: usize,
118    pub candidate_count: usize,
119    pub returned_count: usize,
120    pub context_bytes: usize,
121    pub elapsed_ms: u64,
122}
123
124/// Internal context packing input grouped by structural role.
125#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
126pub struct CodeGraphContextPack {
127    pub entry_points: Vec<CodeRetrievalHit>,
128    pub related_symbols: Vec<CodeRetrievalHit>,
129    pub graph_paths: Vec<CodeRetrievalHit>,
130    pub impact_hints: Vec<CodeGraphImpactHint>,
131    pub code_excerpts: Vec<CodeGraphCodeExcerpt>,
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    #[test]
139    fn context_request_rejects_empty_query() {
140        let error = request(" ", 1, 1024).expect_err("empty query should fail");
141
142        assert!(error.to_string().contains("query"));
143    }
144
145    #[test]
146    fn context_request_bounds_limit_and_context_bytes() {
147        assert!(request("retry", 0, 1024).is_err());
148        assert!(request("retry", CODEGRAPH_CONTEXT_MAX_LIMIT + 1, 1024).is_err());
149        assert!(request("retry", 1, 0).is_err());
150        assert!(request("retry", 1, CODEGRAPH_CONTEXT_MIN_BYTES - 1).is_err());
151        assert!(request("retry", 1, CODEGRAPH_CONTEXT_MAX_BYTES + 1).is_err());
152        assert!(
153            request(
154                "retry",
155                CODEGRAPH_CONTEXT_MAX_LIMIT,
156                CODEGRAPH_CONTEXT_MIN_BYTES
157            )
158            .is_ok()
159        );
160    }
161
162    #[test]
163    fn context_request_defaults_optional_code_toggles_from_json() {
164        let request: CodeGraphContextRequest = serde_json::from_value(serde_json::json!({
165            "repository": {
166                "repository": "repo",
167                "ref_selector": "HEAD",
168                "path_filters": [],
169                "language_filters": []
170            },
171            "query": "retry",
172            "limit": 1,
173            "freshness_policy": "allow_stale",
174            "max_context_bytes": CODEGRAPH_CONTEXT_MIN_BYTES
175        }))
176        .expect("request should deserialize with default toggles");
177
178        assert!(request.include_code);
179        assert!(!request.exclude_generated);
180    }
181
182    fn request(
183        query: &str,
184        limit: usize,
185        max_context_bytes: usize,
186    ) -> Result<CodeGraphContextRequest, DomainError> {
187        CodeGraphContextRequest::new(
188            CodeRepositorySelector::new("repo", "HEAD", Vec::new(), Vec::new())?,
189            query,
190            limit,
191            FreshnessPolicy::AllowStale,
192            max_context_bytes,
193            true,
194            false,
195        )
196    }
197}