From 80ac2bce7dee10aca2eb5d4783e3e47a5fea7cc6 Mon Sep 17 00:00:00 2001
From: sphaerophoria <mick@sayson.com>
Date: Thu, 16 Jan 2025 17:57:30 -0800
Subject: [PATCH] Send $/progress messages on build file parsing (#2141)

---
 src/DocumentStore.zig | 99 +++++++++++++++++++++++++++++++++++++++++++
 src/Server.zig        |  7 +++
 2 files changed, 106 insertions(+)

diff --git a/src/DocumentStore.zig b/src/DocumentStore.zig
index 498c0c6f8..3ec1fda20 100644
--- a/src/DocumentStore.zig
+++ b/src/DocumentStore.zig
@@ -6,6 +6,7 @@ const URI = @import("uri.zig");
 const analysis = @import("analysis.zig");
 const offsets = @import("offsets.zig");
 const log = std.log.scoped(.store);
+const lsp = @import("lsp");
 const Ast = std.zig.Ast;
 const BuildAssociatedConfig = @import("BuildAssociatedConfig.zig");
 const BuildConfig = @import("build_runner/shared.zig").BuildConfig;
@@ -25,6 +26,9 @@ handles: std.StringArrayHashMapUnmanaged(*Handle) = .empty,
 build_files: std.StringArrayHashMapUnmanaged(*BuildFile) = .empty,
 cimports: std.AutoArrayHashMapUnmanaged(Hash, translate_c.Result) = .empty,
 diagnostics_collection: *DiagnosticsCollection,
+builds_in_progress: std.atomic.Value(i32) = .init(0),
+transport: ?lsp.AnyTransport = null,
+supports_work_done_progress: bool = false,
 
 pub const Uri = []const u8;
 
@@ -833,9 +837,101 @@ pub fn invalidateBuildFile(self: *DocumentStore, build_file_uri: Uri) void {
     };
 }
 
+const progress_token = "buildProgressToken";
+
+fn sendMessageToClient(allocator: std.mem.Allocator, transport: lsp.AnyTransport, message: anytype) !void {
+    const serialized = try std.json.stringifyAlloc(
+        allocator,
+        message,
+        .{ .emit_null_optional_fields = false },
+    );
+    defer allocator.free(serialized);
+
+    try transport.writeJsonMessage(serialized);
+}
+
+fn notifyBuildStart(self: *DocumentStore) void {
+    if (!self.supports_work_done_progress) return;
+
+    // Atomicity note: We do not actually care about memory surrounding the
+    // counter, we only care about the counter itself. We only need to ensure
+    // we aren't double entering/exiting
+    const prev = self.builds_in_progress.fetchAdd(1, .monotonic);
+    if (prev != 0) return;
+
+    const transport = self.transport orelse return;
+
+    sendMessageToClient(
+        self.allocator,
+        transport,
+        .{
+            .jsonrpc = "2.0",
+            .id = "progress",
+            .method = "window/workDoneProgress/create",
+            .params = lsp.types.WorkDoneProgressCreateParams{
+                .token = .{ .string = progress_token },
+            },
+        },
+    ) catch |err| {
+        log.err("Failed to send create work message: {}", .{err});
+        return;
+    };
+
+    sendMessageToClient(self.allocator, transport, .{
+        .jsonrpc = "2.0",
+        .method = "$/progress",
+        .params = .{
+            .token = progress_token,
+            .value = lsp.types.WorkDoneProgressBegin{
+                .title = "Loading build configuration",
+            },
+        },
+    }) catch |err| {
+        log.err("Failed to send progress start message: {}", .{err});
+        return;
+    };
+}
+
+const EndStatus = enum { success, failed };
+
+fn notifyBuildEnd(self: *DocumentStore, status: EndStatus) void {
+    if (!self.supports_work_done_progress) return;
+
+    // Atomicity note: We do not actually care about memory surrounding the
+    // counter, we only care about the counter itself. We only need to ensure
+    // we aren't double entering/exiting
+    const prev = self.builds_in_progress.fetchSub(1, .monotonic);
+    if (prev != 1) return;
+
+    const transport = self.transport orelse return;
+
+    const message = switch (status) {
+        .failed => "Failed",
+        .success => "Success",
+    };
+
+    sendMessageToClient(self.allocator, transport, .{
+        .jsonrpc = "2.0",
+        .method = "$/progress",
+        .params = .{
+            .token = progress_token,
+            .value = lsp.types.WorkDoneProgressEnd{
+                .message = message,
+            },
+        },
+    }) catch |err| {
+        log.err("Failed to send progress end message: {}", .{err});
+        return;
+    };
+}
+
 fn invalidateBuildFileWorker(self: *DocumentStore, build_file_uri: Uri, is_build_file_uri_owned: bool) void {
     defer if (is_build_file_uri_owned) self.allocator.free(build_file_uri);
 
+    var end_status: EndStatus = .failed;
+    self.notifyBuildStart();
+    defer self.notifyBuildEnd(end_status);
+
     const build_config = loadBuildConfiguration(self, build_file_uri) catch |err| {
         log.err("Failed to load build configuration for {s} (error: {})", .{ build_file_uri, err });
         return;
@@ -846,6 +942,9 @@ fn invalidateBuildFileWorker(self: *DocumentStore, build_file_uri: Uri, is_build
         return;
     };
     build_file.setBuildConfig(build_config);
+
+    // Looks like a useless assignment, but alters deffered onEnd
+    end_status = .success;
 }
 
 /// The `DocumentStore` represents a graph structure where every
diff --git a/src/Server.zig b/src/Server.zig
index 1ac3025e2..ea1afcb69 100644
--- a/src/Server.zig
+++ b/src/Server.zig
@@ -489,6 +489,12 @@ fn initializeHandler(server: *Server, arena: std.mem.Allocator, request: types.I
         }
     }
 
+    if (request.capabilities.window) |window| {
+        if (window.workDoneProgress) |wdp| {
+            server.document_store.supports_work_done_progress = wdp;
+        }
+    }
+
     if (request.capabilities.workspace) |workspace| {
         server.client_capabilities.supports_apply_edits = workspace.applyEdit orelse false;
         server.client_capabilities.supports_configuration = workspace.configuration orelse false;
@@ -1904,6 +1910,7 @@ pub fn destroy(server: *Server) void {
 pub fn setTransport(server: *Server, transport: lsp.AnyTransport) void {
     server.transport = transport;
     server.diagnostics_collection.transport = transport;
+    server.document_store.transport = transport;
 }
 
 pub fn keepRunning(server: Server) bool {