Skip to main content

ai_session_search/
refs.rs

1// SPDX-FileCopyrightText: 2026 Andrew Hundt
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::{BTreeMap, HashSet};
5
6use linkify::{LinkFinder, LinkKind};
7use serde::Serialize;
8use url::Url;
9
10#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
11pub struct MessageRef {
12    pub kind: String,
13    pub value: String,
14    pub normalized_value: Option<String>,
15    pub host: Option<String>,
16    pub source_tool: Option<String>,
17    pub source_field: Option<String>,
18    pub confidence: String,
19    pub span_start: usize,
20    pub span_end: usize,
21}
22
23pub fn extract_refs_from_text(text: &str, source_tool: Option<&str>) -> Vec<MessageRef> {
24    let mut finder = LinkFinder::new();
25    finder.kinds(&[LinkKind::Url]);
26    finder.url_must_have_scheme(false);
27
28    let mut seen = HashSet::new();
29    let mut refs = Vec::new();
30    for link in finder.links(text) {
31        let raw = trim_trailing_punctuation(link.as_str());
32        if raw.is_empty() {
33            continue;
34        }
35        if is_probable_local_file_ref(raw) {
36            continue;
37        }
38        let span_start = link.start();
39        let span_end = span_start + raw.len();
40        let (normalized_value, host) = normalize_url(raw);
41        let key = (
42            raw.to_ascii_lowercase(),
43            normalized_value.clone(),
44            span_start,
45            span_end,
46        );
47        if !seen.insert(key) {
48            continue;
49        }
50        refs.push(MessageRef {
51            kind: "url".to_string(),
52            value: raw.to_string(),
53            normalized_value,
54            host,
55            source_tool: source_tool.map(str::to_string),
56            source_field: None,
57            confidence: "parsed".to_string(),
58            span_start,
59            span_end,
60        });
61    }
62    refs
63}
64
65pub fn ref_summary(refs: &[MessageRef]) -> String {
66    if refs.is_empty() {
67        return String::new();
68    }
69    let mut counts: BTreeMap<&str, usize> = BTreeMap::new();
70    for item in refs {
71        *counts.entry(item.kind.as_str()).or_default() += 1;
72    }
73    counts
74        .into_iter()
75        .map(|(kind, count)| {
76            if count == 1 {
77                kind.to_string()
78            } else {
79                format!("{count} {kind}s")
80            }
81        })
82        .collect::<Vec<_>>()
83        .join(", ")
84}
85
86fn normalize_url(raw: &str) -> (Option<String>, Option<String>) {
87    let candidate = if raw.contains("://") {
88        raw.to_string()
89    } else {
90        format!("https://{raw}")
91    };
92    match Url::parse(&candidate) {
93        Ok(url) if matches!(url.scheme(), "http" | "https") => {
94            let host = url.host_str().map(str::to_string);
95            (Some(url.to_string()), host)
96        }
97        _ => (None, None),
98    }
99}
100
101fn trim_trailing_punctuation(raw: &str) -> &str {
102    raw.trim_end_matches(['.', ',', ';', ':', ')', ']', '}'])
103}
104
105fn is_probable_local_file_ref(raw: &str) -> bool {
106    let lower = raw.to_ascii_lowercase();
107    if lower.contains("://") || lower.starts_with("www.") {
108        return false;
109    }
110    if lower.contains('/') {
111        return is_probable_local_path_ref(&lower);
112    }
113    let bare = strip_line_suffix(&lower);
114    if matches!(bare, "docs.rs") {
115        return false;
116    }
117    let Some((stem, ext)) = bare.rsplit_once('.') else {
118        return false;
119    };
120    !stem.is_empty() && is_common_file_extension(ext)
121}
122
123fn is_probable_local_path_ref(raw: &str) -> bool {
124    let first = raw.split('/').find(|part| !part.is_empty()).unwrap_or("");
125    if first.contains('.') && !is_probable_local_file_ref(first) {
126        return false;
127    }
128    raw.rsplit('/')
129        .find(|part| !part.is_empty())
130        .is_some_and(is_probable_local_file_ref)
131}
132
133fn strip_line_suffix(value: &str) -> &str {
134    let Some((prefix, suffix)) = value.split_once(':') else {
135        return value;
136    };
137    let line_ref = suffix.split_once('-').map_or_else(
138        || !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()),
139        |(start, end)| {
140            !start.is_empty()
141                && !end.is_empty()
142                && start.chars().all(|c| c.is_ascii_digit())
143                && end.chars().all(|c| c.is_ascii_digit())
144        },
145    );
146    if line_ref {
147        prefix
148    } else {
149        value
150    }
151}
152
153fn is_common_file_extension(ext: &str) -> bool {
154    matches!(
155        ext,
156        "adoc"
157            | "astro"
158            | "awk"
159            | "bash"
160            | "bat"
161            | "bib"
162            | "c"
163            | "cc"
164            | "cfg"
165            | "clj"
166            | "cljs"
167            | "cmake"
168            | "conf"
169            | "cpp"
170            | "cr"
171            | "cs"
172            | "csh"
173            | "css"
174            | "csv"
175            | "dart"
176            | "diff"
177            | "env"
178            | "erb"
179            | "ex"
180            | "exs"
181            | "fish"
182            | "fs"
183            | "gemspec"
184            | "go"
185            | "graphql"
186            | "groovy"
187            | "h"
188            | "haml"
189            | "hbs"
190            | "heex"
191            | "hh"
192            | "hpp"
193            | "hs"
194            | "htm"
195            | "html"
196            | "ini"
197            | "ipynb"
198            | "java"
199            | "js"
200            | "json"
201            | "jsonl"
202            | "jsonnet"
203            | "jsx"
204            | "kt"
205            | "kts"
206            | "less"
207            | "lhs"
208            | "lock"
209            | "lua"
210            | "m"
211            | "md"
212            | "mdx"
213            | "mk"
214            | "mm"
215            | "nix"
216            | "patch"
217            | "pdf"
218            | "php"
219            | "plist"
220            | "proto"
221            | "ps1"
222            | "psm1"
223            | "py"
224            | "pyi"
225            | "r"
226            | "rake"
227            | "rb"
228            | "rst"
229            | "rs"
230            | "scala"
231            | "scss"
232            | "sh"
233            | "sol"
234            | "sql"
235            | "svelte"
236            | "svg"
237            | "swift"
238            | "tex"
239            | "toml"
240            | "twig"
241            | "ts"
242            | "tsx"
243            | "tsv"
244            | "txt"
245            | "vim"
246            | "vue"
247            | "wasm"
248            | "xml"
249            | "yaml"
250            | "yml"
251            | "zsh"
252    )
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn extracts_explicit_and_scheme_less_urls_with_spans() {
261        let text = "See https://example.com/paper.pdf, docs.rs/linkify, and www.rust-lang.org).";
262        let refs = extract_refs_from_text(text, Some("WebSearch"));
263        let values = refs.iter().map(|r| r.value.as_str()).collect::<Vec<_>>();
264        assert!(values.contains(&"https://example.com/paper.pdf"));
265        assert!(values.contains(&"docs.rs/linkify"));
266        assert!(values.contains(&"www.rust-lang.org"));
267        assert!(refs
268            .iter()
269            .all(|r| r.source_tool.as_deref() == Some("WebSearch")));
270        assert!(refs.iter().all(|r| r.span_start < r.span_end));
271        assert!(refs
272            .iter()
273            .any(|r| r.host.as_deref() == Some("example.com")));
274    }
275
276    #[test]
277    fn summarizes_refs_by_kind() {
278        let refs = extract_refs_from_text("https://a.test https://b.test", None);
279        assert_eq!(ref_summary(&refs), "2 urls");
280        assert_eq!(ref_summary(&[]), "");
281    }
282
283    #[test]
284    fn filters_filename_like_scheme_less_refs_without_dropping_domains() {
285        let text = "See util.rs, cli.py:42, db.rs:75-91, paper.tex, refs.bib, notes.md, notebook.ipynb, figure.svg, src/db.rs, docs/refs.bib, docs.rs/linkify, github.com/org/repo, and example.com/path.";
286        let refs = extract_refs_from_text(text, None);
287        let values = refs.iter().map(|r| r.value.as_str()).collect::<Vec<_>>();
288        assert!(!values.contains(&"util.rs"));
289        assert!(!values.contains(&"cli.py:42"));
290        assert!(!values.contains(&"db.rs:75-91"));
291        assert!(!values.contains(&"paper.tex"));
292        assert!(!values.contains(&"refs.bib"));
293        assert!(!values.contains(&"notes.md"));
294        assert!(!values.contains(&"notebook.ipynb"));
295        assert!(!values.contains(&"figure.svg"));
296        assert!(!values.contains(&"src/db.rs"));
297        assert!(!values.contains(&"docs/refs.bib"));
298        assert!(values.contains(&"docs.rs/linkify"));
299        assert!(values.contains(&"github.com/org/repo"));
300        assert!(values.contains(&"example.com/path"));
301    }
302}