tenant.openagents/omega

1use anyhow::{Context, Result};
2use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity};
3use lsp::LanguageServer;
4
5use crate::LspStore;
6
7const LOGGER: zlog::Logger = zlog::scoped!("json-schema");
8
9/// https://github.com/Microsoft/vscode/blob/main/extensions/json-language-features/server/README.md#schema-content-request
10///
11/// Represents a "JSON language server-specific, non-standardized, extension to the LSP" with which the vscode-json-language-server
12/// can request the contents of a schema that is associated with a uri scheme it does not support.
13/// In our case, we provide the uris for actions on server startup under the `zed://schemas/action/{normalize_action_name}` scheme.
14/// We can then respond to this request with the schema content on demand, thereby greatly reducing the total size of the JSON we send to the server on startup
15struct SchemaContentRequest {}
16
17impl lsp::request::Request for SchemaContentRequest {
18    type Params = Vec<String>;
19
20    type Result = String;
21
22    const METHOD: &'static str = "vscode/content";
23}
24
25type SchemaRequestHandler = fn(Entity<LspStore>, String, &mut AsyncApp) -> Task<Result<String>>;
26pub struct SchemaHandlingImpl(SchemaRequestHandler);
27
28impl Global for SchemaHandlingImpl {}
29
30pub fn register_schema_handler(handler: SchemaRequestHandler, cx: &mut App) {
31    debug_assert!(
32        !cx.has_global::<SchemaHandlingImpl>(),
33        "SchemaHandlingImpl already registered"
34    );
35    cx.set_global(SchemaHandlingImpl(handler));
36}
37
38struct SchemaContentsChanged {}
39
40impl lsp::notification::Notification for SchemaContentsChanged {
41    const METHOD: &'static str = "json/schemaContent";
42    type Params = String;
43}
44
45pub fn notify_schemas_changed(lsp_store: Entity<LspStore>, uris: &[String], cx: &App) {
46    zlog::trace!(LOGGER => "Notifying schema changes for URIs: {:?}", uris);
47    let servers = lsp_store.read_with(cx, |lsp_store, _| {
48        let mut servers = Vec::new();
49        let Some(local) = lsp_store.as_local() else {
50            return servers;
51        };
52
53        for states in local.language_servers.values() {
54            let json_server = match states {
55                super::LanguageServerState::Running {
56                    adapter, server, ..
57                } if adapter.adapter.is_primary_zed_json_schema_adapter() => server.clone(),
58                _ => continue,
59            };
60
61            servers.push(json_server);
62        }
63        servers
64    });
65    for server in servers {
66        for uri in uris {
67            zlog::trace!(LOGGER => "Notifying server {NAME} (id {ID:?}) of schema change for URI: {uri:?}",
68                NAME = server.name(),
69                ID = server.server_id()
70            );
71            if let Err(error) = server.notify::<SchemaContentsChanged>(uri.clone()) {
72                zlog::error!(
73                    LOGGER => "Failed to notify server {NAME} (id {ID:?}) of schema change for URI {uri:?}: {error:#}",
74                        NAME = server.name(),
75                        ID = server.server_id(),
76                );
77            }
78        }
79    }
80}
81
82pub fn register_requests(lsp_store: WeakEntity<LspStore>, language_server: &LanguageServer) {
83    language_server
84        .on_request::<SchemaContentRequest, _, _>(move |params, cx| {
85            let handler = cx.try_read_global::<SchemaHandlingImpl, _>(|handler, _| handler.0);
86            let mut cx = cx.clone();
87            let uri = params.clone().pop();
88            let lsp_store = lsp_store.clone();
89            let resolution = async move {
90                let lsp_store = lsp_store.upgrade().context("LSP store has been dropped")?;
91                let uri = uri.context("No URI")?;
92                let handle_schema_request = handler.context("No schema handler registered")?;
93                handle_schema_request(lsp_store, uri, &mut cx).await
94            };
95            async move {
96                zlog::trace!(LOGGER => "Handling schema request for {:?}", &params);
97                let result = resolution.await;
98                match &result {
99                    Ok(content) => {zlog::trace!(LOGGER => "Schema request resolved with {}B schema", content.len());},
100                    Err(err) => {zlog::warn!(LOGGER => "Schema request failed: {}", err);},
101                }
102                result
103            }
104        })
105        .detach();
106}
107