Skip to repository content

tenant.openagents/omega

No repository description is available.

OpenAgents Git authority 2026-07-28T02:52:12.472Z Public web read
NIP-34 coordinate30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omega
MaintainersHidden in public view
References2 branches · 1 tag
Read-only clonegit clone https://openagents.com/git/tenant.openagents/omega.git
Browse files

semantic_tokens.rs

1140 lines · 42.5 KB · rust
1use std::{collections::hash_map, ops::Range, slice::ChunksExact, sync::Arc};
2
3use anyhow::Result;
4
5use clock::Global;
6use collections::{HashMap, HashSet};
7use futures::{
8    FutureExt as _,
9    future::{Shared, join_all},
10};
11use gpui::{App, AppContext, AsyncApp, Context, Entity, ReadGlobal as _, SharedString, Task};
12use language::{Buffer, LanguageName, language_settings::all_language_settings};
13use lsp::{AdapterServerCapabilities, LanguageServerId};
14use rpc::{TypedEnvelope, proto};
15use settings::{
16    DefaultSemanticTokenRules, SemanticTokenRule, SemanticTokenRules, Settings as _, SettingsStore,
17};
18use smol::future::yield_now;
19
20use text::{Anchor, Bias, OffsetUtf16, PointUtf16, Unclipped};
21use util::ResultExt as _;
22
23use crate::{
24    LanguageServerToQuery, LspStore, LspStoreEvent,
25    lsp_command::{
26        LspCommand, SemanticTokensDelta, SemanticTokensEdit, SemanticTokensFull,
27        SemanticTokensResponse,
28    },
29    lsp_store::missing_servers_to_query,
30    project_settings::ProjectSettings,
31};
32
33pub(super) struct SemanticTokenConfig {
34    stylizers: HashMap<(LanguageServerId, Option<LanguageName>), SemanticTokenStylizer>,
35    rules: SemanticTokenRules,
36    global_mode: settings::SemanticTokens,
37}
38
39impl SemanticTokenConfig {
40    pub(super) fn new(cx: &App) -> Self {
41        Self {
42            stylizers: HashMap::default(),
43            rules: ProjectSettings::get_global(cx)
44                .global_lsp_settings
45                .semantic_token_rules
46                .clone(),
47            global_mode: all_language_settings(None, cx).defaults.semantic_tokens,
48        }
49    }
50
51    pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
52        self.stylizers.retain(|&(id, _), _| id != server_id);
53    }
54
55    pub(super) fn update_rules(&mut self, new_rules: SemanticTokenRules) -> bool {
56        if new_rules != self.rules {
57            self.rules = new_rules;
58            self.stylizers.clear();
59            true
60        } else {
61            false
62        }
63    }
64
65    /// Clears all cached stylizers.
66    ///
67    /// This is called when settings change to ensure that any modifications to
68    /// language-specific semantic token rules (e.g. from extension install/uninstall)
69    /// are picked up. Stylizers are recreated lazily on next use.
70    pub(super) fn clear_stylizers(&mut self) {
71        self.stylizers.clear();
72    }
73
74    pub(super) fn update_global_mode(&mut self, new_mode: settings::SemanticTokens) -> bool {
75        if new_mode != self.global_mode {
76            self.global_mode = new_mode;
77            true
78        } else {
79            false
80        }
81    }
82}
83
84impl LspStore {
85    pub fn semantic_tokens(
86        &mut self,
87        buffer: Entity<Buffer>,
88        cx: &mut Context<Self>,
89    ) -> SemanticTokensTask {
90        let version_queried_for = buffer.read(cx).version();
91        let current_servers = self.relevant_server_ids_for_capability_check(&buffer, cx);
92        let latest_lsp_data = self.latest_lsp_data(&buffer, cx);
93        let semantic_tokens_data = latest_lsp_data.semantic_tokens.get_or_insert_default();
94        let refreshed_servers = std::mem::take(&mut semantic_tokens_data.pending_refreshes);
95        if !refreshed_servers.is_empty() {
96            semantic_tokens_data.update = None;
97            semantic_tokens_data.generation += 1;
98        }
99        for refreshed_server in &refreshed_servers {
100            semantic_tokens_data
101                .raw_tokens
102                .servers
103                .remove(refreshed_server);
104        }
105        let missing_servers = missing_servers_to_query(
106            &mut semantic_tokens_data.raw_tokens.servers,
107            &mut semantic_tokens_data.fetched_servers,
108            &current_servers,
109        )
110        .unwrap_or_default();
111        if !missing_servers.is_empty() {
112            semantic_tokens_data.update = None;
113        }
114        let query_generation = semantic_tokens_data.generation;
115
116        if let Some((updating_for, task)) = &semantic_tokens_data.update
117            && !version_queried_for.changed_since(updating_for)
118        {
119            return task.clone();
120        }
121
122        let mut servers_to_fetch = refreshed_servers;
123        servers_to_fetch.extend(missing_servers);
124        let for_server = if servers_to_fetch.len() == 1 {
125            servers_to_fetch.iter().next().copied()
126        } else {
127            // With multiple servers refreshed, query all of them instead of fanning out
128            // filtered queries — the non-refreshed ones kept their tokens and answer
129            // with cheap deltas.
130            None
131        };
132        semantic_tokens_data
133            .fetched_servers
134            .extend(match for_server {
135                Some(server_id) => HashSet::from_iter([server_id]),
136                None => current_servers,
137            });
138        let new_tokens = self.fetch_semantic_tokens_for_buffer(&buffer, for_server, cx);
139
140        let task_buffer = buffer.clone();
141        let task_version_queried_for = version_queried_for.clone();
142        let task = cx
143            .spawn(async move |lsp_store, cx| {
144                let buffer = task_buffer;
145                let version_queried_for = task_version_queried_for;
146                let res = if let Some(new_tokens) = new_tokens.await {
147                    let (raw_tokens, buffer_snapshot) = lsp_store
148                        .update(cx, |lsp_store, cx| {
149                            let lsp_data = lsp_store.latest_lsp_data(&buffer, cx);
150                            let semantic_tokens_data =
151                                lsp_data.semantic_tokens.get_or_insert_default();
152
153                            if version_queried_for == lsp_data.buffer_version
154                                && semantic_tokens_data.generation == query_generation
155                            {
156                                for (server_id, new_tokens_response) in new_tokens {
157                                    match new_tokens_response {
158                                        SemanticTokensResponse::Full { data, result_id } => {
159                                            semantic_tokens_data.raw_tokens.servers.insert(
160                                                server_id,
161                                                Arc::new(ServerSemanticTokens::from_full(
162                                                    data, result_id,
163                                                )),
164                                            );
165                                        }
166                                        SemanticTokensResponse::Delta { edits, result_id } => {
167                                            if let Some(tokens) = semantic_tokens_data
168                                                .raw_tokens
169                                                .servers
170                                                .get_mut(&server_id)
171                                            {
172                                                let tokens = Arc::make_mut(tokens);
173                                                tokens.result_id = result_id;
174                                                tokens.apply(&edits);
175                                            }
176                                        }
177                                    }
178                                }
179                            }
180                            let buffer_snapshot =
181                                buffer.read_with(cx, |buffer, _| buffer.snapshot());
182                            (semantic_tokens_data.raw_tokens.clone(), buffer_snapshot)
183                        })
184                        .map_err(Arc::new)?;
185                    Some(
186                        cx.background_spawn(raw_to_buffer_semantic_tokens(
187                            raw_tokens,
188                            buffer_snapshot.text.clone(),
189                        ))
190                        .await,
191                    )
192                } else {
193                    lsp_store.update(cx, |lsp_store, cx| {
194                        if let Some(current_lsp_data) =
195                            lsp_store.current_lsp_data(buffer.read(cx).remote_id())
196                            && current_lsp_data.buffer_version == version_queried_for
197                            && let Some(semantic_tokens) = current_lsp_data.semantic_tokens.as_mut()
198                            && semantic_tokens.generation == query_generation
199                        {
200                            semantic_tokens.evict_all();
201                        }
202                    })?;
203                    None
204                };
205                Ok(BufferSemanticTokens { tokens: res })
206            })
207            .shared();
208
209        self.latest_lsp_data(&buffer, cx)
210            .semantic_tokens
211            .get_or_insert_default()
212            .update = Some((version_queried_for, task.clone()));
213
214        task
215    }
216
217    pub(super) fn fetch_semantic_tokens_for_buffer(
218        &mut self,
219        buffer: &Entity<Buffer>,
220        for_server: Option<LanguageServerId>,
221        cx: &mut Context<Self>,
222    ) -> Task<Option<HashMap<LanguageServerId, SemanticTokensResponse>>> {
223        if let Some((client, upstream_project_id)) = self.upstream_client() {
224            let request = SemanticTokensFull { for_server };
225            if !self.is_capable_for_proto_request(buffer, &request, cx) {
226                return Task::ready(None);
227            }
228
229            let request_timeout = ProjectSettings::get_global(cx)
230                .global_lsp_settings
231                .get_request_timeout();
232            let request_task = client.request_lsp(
233                upstream_project_id,
234                None,
235                request_timeout,
236                cx.background_executor().clone(),
237                request.to_proto(upstream_project_id, buffer.read(cx)),
238            );
239            let buffer = buffer.clone();
240            cx.spawn(async move |weak_lsp_store, cx| {
241                let lsp_store = weak_lsp_store.upgrade()?;
242                let tokens = join_all(
243                    request_task
244                        .await
245                        .log_err()
246                        .flatten()
247                        .map(|response| response.payload)
248                        .unwrap_or_default()
249                        .into_iter()
250                        .map(|response| {
251                            let server_id = LanguageServerId::from_proto(response.server_id);
252                            let response = request.response_from_proto(
253                                response.response,
254                                lsp_store.clone(),
255                                buffer.clone(),
256                                cx.clone(),
257                            );
258                            async move {
259                                match response.await {
260                                    Ok(tokens) => Some((server_id, tokens)),
261                                    Err(e) => {
262                                        log::error!("Failed to query remote semantic tokens for server {server_id:?}: {e:#}");
263                                        None
264                                    }
265                                }
266                            }
267                        }),
268                )
269                .await
270                .into_iter()
271                .flatten()
272                .collect();
273                Some(tokens)
274            })
275        } else {
276            let token_tasks = self
277                .local_lsp_servers_for_buffer(&buffer, cx)
278                .into_iter()
279                .filter(|&server_id| {
280                    for_server.is_none_or(|for_server_id| for_server_id == server_id)
281                })
282                .filter_map(|server_id| {
283                    let capabilities = AdapterServerCapabilities {
284                        server_capabilities: self.lsp_server_capabilities.get(&server_id)?.clone(),
285                        code_action_kinds: None,
286                    };
287                    let request_task = match self.semantic_tokens_result_id(server_id, buffer, cx) {
288                        Some(result_id) => {
289                            let delta_request = SemanticTokensDelta {
290                                previous_result_id: result_id,
291                            };
292                            if !delta_request.check_capabilities(capabilities.clone()) {
293                                let full_request = SemanticTokensFull {
294                                    for_server: Some(server_id),
295                                };
296                                if !full_request.check_capabilities(capabilities) {
297                                    return None;
298                                }
299
300                                self.request_lsp(
301                                    buffer.clone(),
302                                    LanguageServerToQuery::Other(server_id),
303                                    full_request,
304                                    cx,
305                                )
306                            } else {
307                                self.request_lsp(
308                                    buffer.clone(),
309                                    LanguageServerToQuery::Other(server_id),
310                                    delta_request,
311                                    cx,
312                                )
313                            }
314                        }
315                        None => {
316                            let request = SemanticTokensFull {
317                                for_server: Some(server_id),
318                            };
319                            if !request.check_capabilities(capabilities) {
320                                return None;
321                            }
322                            self.request_lsp(
323                                buffer.clone(),
324                                LanguageServerToQuery::Other(server_id),
325                                request,
326                                cx,
327                            )
328                        }
329                    };
330                    Some(async move { (server_id, request_task.await) })
331                })
332                .collect::<Vec<_>>();
333            if token_tasks.is_empty() {
334                return Task::ready(None);
335            }
336
337            cx.background_spawn(async move {
338                Some(
339                    join_all(token_tasks)
340                        .await
341                        .into_iter()
342                        .flat_map(|(server_id, response)| {
343                            match response {
344                                Ok(tokens) => Some((server_id, tokens)),
345                                Err(e) => {
346                                    log::error!("Failed to query remote semantic tokens for server {server_id:?}: {e:#}");
347                                    None
348                                }
349                            }
350                        })
351                        .collect()
352                )
353            })
354        }
355    }
356
357    /// Marks the server's semantic tokens as refresh-pending in every buffer, to be
358    /// invalidated by the next query, and notifies the observers.
359    pub(crate) fn refresh_semantic_tokens(
360        &mut self,
361        server_id: LanguageServerId,
362        cx: &mut Context<Self>,
363    ) {
364        for lsp_data in self.lsp_data.values_mut() {
365            if let Some(semantic_tokens) = &mut lsp_data.semantic_tokens {
366                semantic_tokens.pending_refreshes.insert(server_id);
367            }
368        }
369        cx.emit(LspStoreEvent::RefreshSemanticTokens { server_id });
370        if let Some((client, project_id)) = self.downstream_client.as_ref() {
371            client
372                .send(proto::RefreshSemanticTokens {
373                    project_id: *project_id,
374                    server_id: server_id.to_proto(),
375                    request_id: Some(super::next_wire_refresh_request_id()),
376                })
377                .log_err();
378        }
379    }
380
381    pub(crate) async fn handle_refresh_semantic_tokens(
382        lsp_store: Entity<Self>,
383        envelope: TypedEnvelope<proto::RefreshSemanticTokens>,
384        mut cx: AsyncApp,
385    ) -> Result<proto::Ack> {
386        lsp_store.update(&mut cx, |lsp_store, cx| {
387            lsp_store.refresh_semantic_tokens(
388                LanguageServerId::from_proto(envelope.payload.server_id),
389                cx,
390            );
391        });
392        Ok(proto::Ack {})
393    }
394
395    #[cfg(any(test, feature = "test-support"))]
396    pub fn semantic_token_servers(&self, buffer_id: text::BufferId) -> Vec<LanguageServerId> {
397        let mut servers = self
398            .lsp_data
399            .get(&buffer_id)
400            .and_then(|lsp_data| lsp_data.semantic_tokens.as_ref())
401            .map(|semantic_tokens| {
402                semantic_tokens
403                    .raw_tokens
404                    .servers
405                    .keys()
406                    .copied()
407                    .collect::<Vec<_>>()
408            })
409            .unwrap_or_default();
410        servers.sort();
411        servers
412    }
413
414    #[cfg(any(test, feature = "test-support"))]
415    pub fn semantic_token_data(
416        &self,
417        buffer_id: text::BufferId,
418    ) -> Vec<(LanguageServerId, Vec<u32>)> {
419        let mut data = self
420            .lsp_data
421            .get(&buffer_id)
422            .and_then(|lsp_data| lsp_data.semantic_tokens.as_ref())
423            .map(|semantic_tokens| {
424                semantic_tokens
425                    .raw_tokens
426                    .servers
427                    .iter()
428                    .map(|(server_id, tokens)| (*server_id, tokens.data.clone()))
429                    .collect::<Vec<_>>()
430            })
431            .unwrap_or_default();
432        data.sort_by_key(|(server_id, _)| *server_id);
433        data
434    }
435
436    fn semantic_tokens_result_id(
437        &mut self,
438        server_id: LanguageServerId,
439        buffer: &Entity<Buffer>,
440        cx: &mut App,
441    ) -> Option<SharedString> {
442        self.latest_lsp_data(buffer, cx)
443            .semantic_tokens
444            .as_ref()?
445            .raw_tokens
446            .servers
447            .get(&server_id)?
448            .result_id
449            .clone()
450    }
451
452    pub fn get_or_create_token_stylizer(
453        &mut self,
454        server_id: LanguageServerId,
455        language: Option<&LanguageName>,
456        cx: &mut App,
457    ) -> Option<&SemanticTokenStylizer> {
458        let stylizer = match self
459            .semantic_token_config
460            .stylizers
461            .entry((server_id, language.cloned()))
462        {
463            hash_map::Entry::Occupied(o) => o.into_mut(),
464            hash_map::Entry::Vacant(v) => {
465                let tokens_provider = self
466                    .lsp_server_capabilities
467                    .get(&server_id)?
468                    .semantic_tokens_provider
469                    .as_ref()?;
470                let legend = match tokens_provider {
471                    lsp::SemanticTokensServerCapabilities::SemanticTokensOptions(opts) => {
472                        &opts.legend
473                    }
474                    lsp::SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
475                        opts,
476                    ) => &opts.semantic_tokens_options.legend,
477                };
478                let language_rules = language.and_then(|language| {
479                    SettingsStore::global(cx).language_semantic_token_rules(language.as_ref())
480                });
481                let stylizer = SemanticTokenStylizer::new(server_id, legend, language_rules, cx);
482                v.insert(stylizer)
483            }
484        };
485        Some(stylizer)
486    }
487}
488
489pub type SemanticTokensTask =
490    Shared<Task<std::result::Result<BufferSemanticTokens, Arc<anyhow::Error>>>>;
491
492#[derive(Debug, Default, Clone)]
493pub struct BufferSemanticTokens {
494    pub tokens: Option<HashMap<LanguageServerId, Arc<[BufferSemanticToken]>>>,
495}
496
497#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
498pub struct TokenType(pub u32);
499
500#[derive(Debug, Clone)]
501pub struct BufferSemanticToken {
502    /// The range of the token in the buffer.
503    ///
504    /// Guaranteed to contain a buffer id.
505    pub range: Range<Anchor>,
506    pub token_type: TokenType,
507    pub token_modifiers: u32,
508}
509
510pub struct SemanticTokenStylizer {
511    server_id: LanguageServerId,
512    rules_by_token_type: HashMap<TokenType, Vec<SemanticTokenRule>>,
513    token_type_names: HashMap<TokenType, SharedString>,
514    modifier_mask: HashMap<SharedString, u32>,
515}
516
517impl SemanticTokenStylizer {
518    pub fn new(
519        server_id: LanguageServerId,
520        legend: &lsp::SemanticTokensLegend,
521        language_rules: Option<&SemanticTokenRules>,
522        cx: &App,
523    ) -> Self {
524        let token_types: HashMap<TokenType, SharedString> = legend
525            .token_types
526            .iter()
527            .enumerate()
528            .map(|(i, token_type)| {
529                (
530                    TokenType(i as u32),
531                    SharedString::from(token_type.as_str().to_string()),
532                )
533            })
534            .collect();
535        let modifier_mask: HashMap<SharedString, u32> = legend
536            .token_modifiers
537            .iter()
538            .enumerate()
539            .map(|(i, modifier)| (SharedString::from(modifier.as_str().to_string()), 1 << i))
540            .collect();
541
542        let global_rules = &ProjectSettings::get_global(cx)
543            .global_lsp_settings
544            .semantic_token_rules;
545        let default_rules = cx.global::<DefaultSemanticTokenRules>();
546
547        let rules_by_token_type = token_types
548            .iter()
549            .map(|(index, token_type_name)| {
550                let filter = |rule: &&SemanticTokenRule| {
551                    rule.token_type
552                        .as_ref()
553                        .is_none_or(|rule_token_type| rule_token_type == token_type_name.as_ref())
554                };
555                let matching_rules: Vec<SemanticTokenRule> = global_rules
556                    .rules
557                    .iter()
558                    .chain(language_rules.into_iter().flat_map(|lr| &lr.rules))
559                    .chain(default_rules.0.rules.iter())
560                    .rev()
561                    .filter(filter)
562                    .cloned()
563                    .collect();
564                (*index, matching_rules)
565            })
566            .collect();
567
568        SemanticTokenStylizer {
569            server_id,
570            rules_by_token_type,
571            token_type_names: token_types,
572            modifier_mask,
573        }
574    }
575
576    pub fn server_id(&self) -> LanguageServerId {
577        self.server_id
578    }
579
580    pub fn token_type_name(&self, token_type: TokenType) -> Option<&SharedString> {
581        self.token_type_names.get(&token_type)
582    }
583
584    pub fn has_modifier(&self, token_modifiers: u32, modifier: &str) -> bool {
585        let Some(mask) = self.modifier_mask.get(modifier) else {
586            return false;
587        };
588        (token_modifiers & mask) != 0
589    }
590
591    pub fn token_modifiers(&self, token_modifiers: u32) -> Option<String> {
592        let modifiers: Vec<&str> = self
593            .modifier_mask
594            .iter()
595            .filter(|(_, mask)| (token_modifiers & *mask) != 0)
596            .map(|(name, _)| name.as_ref())
597            .collect();
598        if modifiers.is_empty() {
599            None
600        } else {
601            Some(modifiers.join(", "))
602        }
603    }
604
605    pub fn rules_for_token(&self, token_type: TokenType) -> Option<&[SemanticTokenRule]> {
606        self.rules_by_token_type
607            .get(&token_type)
608            .map(|v| v.as_slice())
609    }
610}
611
612async fn raw_to_buffer_semantic_tokens(
613    raw_tokens: RawSemanticTokens,
614    buffer_snapshot: text::BufferSnapshot,
615) -> HashMap<LanguageServerId, Arc<[BufferSemanticToken]>> {
616    let mut res = HashMap::default();
617    for (&server_id, server_tokens) in &raw_tokens.servers {
618        let mut last = 0;
619        // We don't do `collect` here due to the filter map not pre-allocating
620        // we'd rather over allocate here than not since we have to re-allocate into an arc slice anyways
621        let mut buffer_tokens = Vec::with_capacity(server_tokens.data.len() / 5);
622        let mut tokens = server_tokens.tokens();
623        // 5000 was chosen by profiling, on a decent machine this will take about 1ms per chunk
624        // This is to avoid blocking the main thread for hundreds of milliseconds at a time for very big files
625        // If we every change the below code to not query the underlying rope 6 times per token we can bump this up
626        const CHUNK_LEN: usize = 5000;
627        loop {
628            let mut changed = false;
629            let chunk = tokens
630                .by_ref()
631                .take(CHUNK_LEN)
632                .inspect(|_| changed = true)
633                .filter_map(|token| {
634                    let start = Unclipped(PointUtf16::new(token.line, token.start));
635                    let clipped_start = buffer_snapshot.clip_point_utf16(start, Bias::Left);
636                    let start_offset = buffer_snapshot
637                        .as_rope()
638                        .point_utf16_to_offset_utf16(clipped_start);
639                    let end_offset = start_offset + OffsetUtf16(token.length as usize);
640
641                    let start = buffer_snapshot
642                        .as_rope()
643                        .offset_utf16_to_offset(start_offset);
644                    if start < last {
645                        return None;
646                    }
647
648                    let end = buffer_snapshot.as_rope().offset_utf16_to_offset(end_offset);
649                    last = end;
650
651                    if start == end {
652                        return None;
653                    }
654
655                    Some(BufferSemanticToken {
656                        range: buffer_snapshot.anchor_range_inside(start..end),
657                        token_type: token.token_type,
658                        token_modifiers: token.token_modifiers,
659                    })
660                });
661            buffer_tokens.extend(chunk);
662
663            if !changed {
664                break;
665            }
666            yield_now().await;
667        }
668
669        res.insert(server_id, buffer_tokens.into());
670    }
671    res
672}
673
674#[derive(Default, Debug)]
675pub struct SemanticTokensData {
676    pub(super) raw_tokens: RawSemanticTokens,
677    pub(super) pending_refreshes: HashSet<LanguageServerId>,
678    fetched_servers: HashSet<LanguageServerId>,
679    update: Option<(Global, SemanticTokensTask)>,
680    /// Bumped on every eviction so that fetches started against the evicted state
681    /// cannot write their stale results into the new one.
682    generation: u64,
683}
684
685impl SemanticTokensData {
686    pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
687        self.raw_tokens.servers.remove(&server_id);
688        self.pending_refreshes.remove(&server_id);
689        self.fetched_servers.remove(&server_id);
690        self.update = None;
691        self.generation += 1;
692    }
693
694    fn evict_all(&mut self) {
695        self.raw_tokens.servers.clear();
696        self.fetched_servers.clear();
697        self.update = None;
698        self.generation += 1;
699    }
700}
701
702/// All the semantic token tokens for a buffer.
703///
704/// This aggregates semantic tokens from multiple language servers in a specific order.
705/// Semantic tokens later in the list will override earlier ones in case of overlap.
706#[derive(Default, Debug, Clone)]
707pub(super) struct RawSemanticTokens {
708    pub servers: HashMap<lsp::LanguageServerId, Arc<ServerSemanticTokens>>,
709}
710
711/// All the semantic tokens for a buffer, from a single language server.
712#[derive(Debug, Clone)]
713pub struct ServerSemanticTokens {
714    /// Each value is:
715    /// data[5*i] - deltaLine: token line number, relative to the start of the previous token
716    /// data[5*i+1] - deltaStart: token start character, relative to the start of the previous token (relative to 0 or the previous token’s start if they are on the same line)
717    /// data[5*i+2] - length: the length of the token.
718    /// data[5*i+3] - tokenType: will be looked up in SemanticTokensLegend.tokenTypes. We currently ask that tokenType < 65536.
719    /// data[5*i+4] - tokenModifiers: each set bit will be looked up in SemanticTokensLegend.tokenModifiers
720    ///
721    /// See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/ for more.
722    data: Vec<u32>,
723
724    pub(crate) result_id: Option<SharedString>,
725}
726
727pub struct SemanticTokensIter<'a> {
728    prev: Option<(u32, u32)>,
729    data: ChunksExact<'a, u32>,
730}
731
732// A single item from `data`.
733struct SemanticTokenValue {
734    delta_line: u32,
735    delta_start: u32,
736    length: u32,
737    token_type: TokenType,
738    token_modifiers: u32,
739}
740
741/// A semantic token, independent of its position.
742#[derive(Debug, PartialEq, Eq)]
743pub struct SemanticToken {
744    pub line: u32,
745    pub start: u32,
746    pub length: u32,
747    pub token_type: TokenType,
748    pub token_modifiers: u32,
749}
750
751impl ServerSemanticTokens {
752    pub fn from_full(data: Vec<u32>, result_id: Option<SharedString>) -> Self {
753        ServerSemanticTokens { data, result_id }
754    }
755
756    pub(crate) fn apply(&mut self, edits: &[SemanticTokensEdit]) {
757        for edit in edits {
758            let start = (edit.start as usize).min(self.data.len());
759            let end = (start + edit.delete_count as usize).min(self.data.len());
760            self.data.splice(start..end, edit.data.iter().copied());
761        }
762    }
763
764    pub fn tokens(&self) -> SemanticTokensIter<'_> {
765        SemanticTokensIter {
766            prev: None,
767            data: self.data.chunks_exact(5),
768        }
769    }
770}
771
772impl Iterator for SemanticTokensIter<'_> {
773    type Item = SemanticToken;
774
775    fn next(&mut self) -> Option<Self::Item> {
776        let chunk = self.data.next()?;
777        let token = SemanticTokenValue {
778            delta_line: chunk[0],
779            delta_start: chunk[1],
780            length: chunk[2],
781            token_type: TokenType(chunk[3]),
782            token_modifiers: chunk[4],
783        };
784
785        let (line, start) = if let Some((last_line, last_start)) = self.prev {
786            let line = last_line + token.delta_line;
787            let start = if token.delta_line == 0 {
788                last_start + token.delta_start
789            } else {
790                token.delta_start
791            };
792            (line, start)
793        } else {
794            (token.delta_line, token.delta_start)
795        };
796
797        self.prev = Some((line, start));
798
799        Some(SemanticToken {
800            line,
801            start,
802            length: token.length,
803            token_type: token.token_type,
804            token_modifiers: token.token_modifiers,
805        })
806    }
807}
808
809#[cfg(test)]
810mod tests {
811    use super::*;
812    use crate::lsp_command::SemanticTokensEdit;
813    use lsp::SEMANTIC_TOKEN_MODIFIERS;
814
815    fn modifier_names(bits: u32) -> String {
816        if bits == 0 {
817            return "-".to_string();
818        }
819        let names: Vec<&str> = SEMANTIC_TOKEN_MODIFIERS
820            .iter()
821            .enumerate()
822            .filter(|(i, _)| bits & (1 << i) != 0)
823            .map(|(_, m)| m.as_str())
824            .collect();
825
826        // Check for unknown bits
827        let known_bits = (1u32 << SEMANTIC_TOKEN_MODIFIERS.len()) - 1;
828        let unknown = bits & !known_bits;
829
830        if unknown != 0 {
831            let mut result = names.join("+");
832            if !result.is_empty() {
833                result.push('+');
834            }
835            result.push_str(&format!("?0x{:x}", unknown));
836            result
837        } else {
838            names.join("+")
839        }
840    }
841
842    /// Debug tool: parses semantic token JSON from LSP and prints human-readable output.
843    ///
844    /// Usage: Paste JSON into `json_input`, then run:
845    ///   cargo test -p project debug_parse_tokens -- --nocapture --ignored
846    ///
847    /// Accepts either:
848    /// - Full LSP response: `{"jsonrpc":"2.0","id":1,"result":{"data":[...]}}`
849    /// - Just the data array: `[0,0,5,1,0,...]`
850    ///
851    /// For delta responses, paste multiple JSON messages (one per line) and they
852    /// will be applied in sequence.
853    ///
854    /// Token encoding (5 values per token):
855    ///   [deltaLine, deltaStart, length, tokenType, tokenModifiers]
856    #[test]
857    #[ignore] // Run with: cargo test -p project debug_parse_tokens -- --nocapture --ignored
858    fn debug_parse_tokens() {
859        // ============================================================
860        // PASTE YOUR JSON HERE (one message per line for sequences)
861        // Comments starting with // are ignored
862        // ============================================================
863        let json_input = r#"
864// === EXAMPLE 1: Full response (LSP spec example) ===
865// 3 tokens: property at line 2, type at line 2, class at line 5
866{"jsonrpc":"2.0","id":1,"result":{"resultId":"1","data":[2,5,3,9,3,0,5,4,6,0,3,2,7,1,0]}}
867
868// === EXAMPLE 2: Delta response ===
869// User added empty line at start of file, so all tokens shift down by 1 line.
870// This changes first token's deltaLine from 2 to 3 (edit at index 0).
871{"jsonrpc":"2.0","id":2,"result":{"resultId":"2","edits":[{"start":0,"deleteCount":1,"data":[3]}]}}
872
873// === EXAMPLE 3: Another delta ===
874// User added a new token. Insert 5 values at position 5 (after first token).
875// New token: same line as token 1, 2 chars after it ends, len 5, type=function(12), mods=definition(2)
876{"jsonrpc":"2.0","id":3,"result":{"resultId":"3","edits":[{"start":5,"deleteCount":0,"data":[0,2,5,12,2]}]}}
877        "#;
878        // Accepted formats:
879        // - Full response: {"result":{"data":[...]}}
880        // - Delta response: {"result":{"edits":[{"start":N,"deleteCount":N,"data":[...]}]}}
881        // - Just array: [0,0,5,1,0,...]
882
883        // ============================================================
884        // PROCESSING
885        // ============================================================
886        let mut current_data: Vec<u32> = Vec::new();
887        let mut result_id: Option<String> = None;
888
889        for line in json_input.lines() {
890            let line = line.trim();
891            if line.is_empty() || line.starts_with("//") {
892                continue;
893            }
894
895            let parsed: serde_json::Value =
896                serde_json::from_str(line).expect("Failed to parse JSON");
897
898            // Try to extract data from various JSON shapes
899            let (data, edits, new_result_id) = extract_semantic_tokens(&parsed);
900
901            if let Some(new_id) = new_result_id {
902                result_id = Some(new_id);
903            }
904
905            if let Some(full_data) = data {
906                println!("\n{}", "=".repeat(70));
907                println!("FULL RESPONSE (resultId: {:?})", result_id);
908                current_data = full_data;
909            } else if let Some(delta_edits) = edits {
910                println!("\n{}", "=".repeat(70));
911                println!(
912                    "DELTA RESPONSE: {} edit(s) (resultId: {:?})",
913                    delta_edits.len(),
914                    result_id
915                );
916                for (i, edit) in delta_edits.iter().enumerate() {
917                    println!(
918                        "  [{}] start={}, delete={}, insert {} values",
919                        i,
920                        edit.start,
921                        edit.delete_count,
922                        edit.data.len()
923                    );
924                }
925                let mut tokens = ServerSemanticTokens::from_full(current_data.clone(), None);
926                tokens.apply(&delta_edits);
927                current_data = tokens.data;
928            }
929        }
930
931        // Print parsed tokens
932        println!(
933            "\nDATA: {} values = {} tokens",
934            current_data.len(),
935            current_data.len() / 5
936        );
937        println!("\nPARSED TOKENS:");
938        println!("{:-<100}", "");
939        println!(
940            "{:>5} {:>6} {:>4}  {:<15} {}",
941            "LINE", "START", "LEN", "TYPE", "MODIFIERS"
942        );
943        println!("{:-<100}", "");
944
945        let tokens = ServerSemanticTokens::from_full(current_data, None);
946        for token in tokens.tokens() {
947            println!(
948                "{:>5} {:>6} {:>4}  {:<15} {}",
949                token.line,
950                token.start,
951                token.length,
952                token.token_type.0,
953                modifier_names(token.token_modifiers),
954            );
955        }
956        println!("{:-<100}", "");
957        println!("{}\n", "=".repeat(100));
958    }
959
960    fn extract_semantic_tokens(
961        value: &serde_json::Value,
962    ) -> (
963        Option<Vec<u32>>,
964        Option<Vec<SemanticTokensEdit>>,
965        Option<String>,
966    ) {
967        // Try as array directly: [1,2,3,...]
968        if let Some(arr) = value.as_array() {
969            let data: Vec<u32> = arr
970                .iter()
971                .filter_map(|v| v.as_u64().map(|n| n as u32))
972                .collect();
973            return (Some(data), None, None);
974        }
975
976        // Try as LSP response: {"result": {"data": [...]} } or {"result": {"edits": [...]}}
977        let result = value.get("result").unwrap_or(value);
978        let result_id = result
979            .get("resultId")
980            .and_then(|v| v.as_str())
981            .map(String::from);
982
983        // Full response with data
984        if let Some(data_arr) = result.get("data").and_then(|v| v.as_array()) {
985            let data: Vec<u32> = data_arr
986                .iter()
987                .filter_map(|v| v.as_u64().map(|n| n as u32))
988                .collect();
989            return (Some(data), None, result_id);
990        }
991
992        // Delta response with edits
993        if let Some(edits_arr) = result.get("edits").and_then(|v| v.as_array()) {
994            let edits: Vec<SemanticTokensEdit> = edits_arr
995                .iter()
996                .filter_map(|e| {
997                    Some(SemanticTokensEdit {
998                        start: e.get("start")?.as_u64()? as u32,
999                        delete_count: e.get("deleteCount")?.as_u64()? as u32,
1000                        data: e
1001                            .get("data")
1002                            .and_then(|d| d.as_array())
1003                            .map(|arr| {
1004                                arr.iter()
1005                                    .filter_map(|v| v.as_u64().map(|n| n as u32))
1006                                    .collect()
1007                            })
1008                            .unwrap_or_default(),
1009                    })
1010                })
1011                .collect();
1012            return (None, Some(edits), result_id);
1013        }
1014
1015        (None, None, result_id)
1016    }
1017
1018    #[test]
1019    fn parses_sample_tokens() {
1020        // Example from the spec: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens
1021        let tokens = ServerSemanticTokens::from_full(
1022            vec![2, 5, 3, 0, 3, 0, 5, 4, 1, 0, 3, 2, 7, 2, 0],
1023            None,
1024        )
1025        .tokens()
1026        .collect::<Vec<SemanticToken>>();
1027
1028        // The spec uses 1-based line numbers, and 0-based character numbers. This test uses 0-based for both.
1029        assert_eq!(
1030            tokens,
1031            &[
1032                SemanticToken {
1033                    line: 2,
1034                    start: 5,
1035                    length: 3,
1036                    token_type: TokenType(0),
1037                    token_modifiers: 3
1038                },
1039                SemanticToken {
1040                    line: 2,
1041                    start: 10,
1042                    length: 4,
1043                    token_type: TokenType(1),
1044                    token_modifiers: 0
1045                },
1046                SemanticToken {
1047                    line: 5,
1048                    start: 2,
1049                    length: 7,
1050                    token_type: TokenType(2),
1051                    token_modifiers: 0
1052                }
1053            ]
1054        );
1055    }
1056
1057    #[test]
1058    fn applies_delta_edit() {
1059        // Example from the spec: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_semanticTokens
1060        // After a user types a new empty line at the beginning of the file,
1061        // the tokens shift down by one line. The delta edit transforms
1062        // [2,5,3,0,3, 0,5,4,1,0, 3,2,7,2,0] into [3,5,3,0,3, 0,5,4,1,0, 3,2,7,2,0]
1063        // by replacing the first element (deltaLine of first token) from 2 to 3.
1064
1065        let mut tokens = ServerSemanticTokens::from_full(
1066            vec![2, 5, 3, 0, 3, 0, 5, 4, 1, 0, 3, 2, 7, 2, 0],
1067            None,
1068        );
1069
1070        tokens.apply(&[SemanticTokensEdit {
1071            start: 0,
1072            delete_count: 1,
1073            data: vec![3],
1074        }]);
1075
1076        let result = tokens.tokens().collect::<Vec<SemanticToken>>();
1077
1078        assert_eq!(
1079            result,
1080            &[
1081                SemanticToken {
1082                    line: 3,
1083                    start: 5,
1084                    length: 3,
1085                    token_type: TokenType(0),
1086                    token_modifiers: 3
1087                },
1088                SemanticToken {
1089                    line: 3,
1090                    start: 10,
1091                    length: 4,
1092                    token_type: TokenType(1),
1093                    token_modifiers: 0
1094                },
1095                SemanticToken {
1096                    line: 6,
1097                    start: 2,
1098                    length: 7,
1099                    token_type: TokenType(2),
1100                    token_modifiers: 0
1101                }
1102            ]
1103        );
1104    }
1105
1106    #[test]
1107    fn applies_out_of_bounds_delta_edit_without_panic() {
1108        let mut tokens = ServerSemanticTokens::from_full(vec![2, 5, 3, 0, 3, 0, 5, 4, 1, 0], None);
1109
1110        // start beyond data length
1111        tokens.apply(&[SemanticTokensEdit {
1112            start: 100,
1113            delete_count: 5,
1114            data: vec![1, 2, 3, 4, 5],
1115        }]);
1116        assert_eq!(
1117            tokens.data,
1118            vec![2, 5, 3, 0, 3, 0, 5, 4, 1, 0, 1, 2, 3, 4, 5]
1119        );
1120
1121        // delete_count extends past data length
1122        let mut tokens = ServerSemanticTokens::from_full(vec![2, 5, 3, 0, 3], None);
1123        tokens.apply(&[SemanticTokensEdit {
1124            start: 3,
1125            delete_count: 100,
1126            data: vec![9, 9],
1127        }]);
1128        assert_eq!(tokens.data, vec![2, 5, 3, 9, 9]);
1129
1130        // empty data
1131        let mut tokens = ServerSemanticTokens::from_full(Vec::new(), None);
1132        tokens.apply(&[SemanticTokensEdit {
1133            start: 0,
1134            delete_count: 5,
1135            data: vec![1, 2, 3, 4, 5],
1136        }]);
1137        assert_eq!(tokens.data, vec![1, 2, 3, 4, 5]);
1138    }
1139}
1140
Served at tenant.openagents/omega Member data and write actions are omitted.