3 files changed, 49 insertions, 0 deletions
diff --git a/src/buildtool/serve_api/remote/serve_api.hpp b/src/buildtool/serve_api/remote/serve_api.hpp
index 50fdc4e3..965cf8df 100644
--- a/src/buildtool/serve_api/remote/serve_api.hpp
+++ b/src/buildtool/serve_api/remote/serve_api.hpp
@@ -58,6 +58,13 @@ class ServeApi final {
             content, archive_type, subdir, resolve_symlinks, sync_tree);
     }
 
+    [[nodiscard]] static auto RetrieveTreeFromDistdir(
+        std::shared_ptr<std::unordered_map<std::string, std::string>> const&
+            distfiles,
+        bool sync_tree = false) -> std::optional<std::string> {
+        return Instance().stc_->ServeDistdirTree(distfiles, sync_tree);
+    }
+
     [[nodiscard]] static auto ContentInRemoteCAS(std::string const& content)
         -> bool {
         return Instance().stc_->ServeContent(content);
diff --git a/src/buildtool/serve_api/remote/source_tree_client.cpp b/src/buildtool/serve_api/remote/source_tree_client.cpp
index 8f4d4960..439612bd 100644
--- a/src/buildtool/serve_api/remote/source_tree_client.cpp
+++ b/src/buildtool/serve_api/remote/source_tree_client.cpp
@@ -121,6 +121,36 @@ auto SourceTreeClient::ServeArchiveTree(
     return response.tree();
 }
 
+auto SourceTreeClient::ServeDistdirTree(
+    std::shared_ptr<std::unordered_map<std::string, std::string>> const&
+        distfiles,
+    bool sync_tree) -> std::optional<std::string> {
+    justbuild::just_serve::ServeDistdirTreeRequest request{};
+    for (auto const& [k, v] : *distfiles) {
+        auto* distfile = request.add_distfiles();
+        distfile->set_name(k);
+        distfile->set_content(v);
+    }
+    request.set_sync_tree(sync_tree);
+
+    grpc::ClientContext context;
+    justbuild::just_serve::ServeDistdirTreeResponse response;
+    grpc::Status status = stub_->ServeDistdirTree(&context, request, &response);
+
+    if (not status.ok()) {
+        LogStatus(&logger_, LogLevel::Debug, status);
+        return std::nullopt;
+    }
+    if (response.status() !=
+        ::justbuild::just_serve::ServeDistdirTreeResponse::OK) {
+        logger_.Emit(LogLevel::Debug,
+                     "ServeDistdirTree response returned with {}",
+                     static_cast<int>(response.status()));
+        return std::nullopt;
+    }
+    return response.tree();
+}
+
 auto SourceTreeClient::ServeContent(std::string const& content) -> bool {
     justbuild::just_serve::ServeContentRequest request{};
     request.set_content(content);
diff --git a/src/buildtool/serve_api/remote/source_tree_client.hpp b/src/buildtool/serve_api/remote/source_tree_client.hpp
index a9b28807..a76d4135 100644
--- a/src/buildtool/serve_api/remote/source_tree_client.hpp
+++ b/src/buildtool/serve_api/remote/source_tree_client.hpp
@@ -54,6 +54,18 @@ class SourceTreeClient {
         std::optional<PragmaSpecial> const& resolve_symlinks,
         bool sync_tree) -> std::optional<std::string>;
 
+    /// \brief Retrieve the Git tree of a directory of distfiles, if all the
+    /// content blobs are known by serve remote.
+    /// \param[in] distfiles Mapping from distfile names to content blob ids.
+    /// \param[in] sync_tree Sync tree and all ditfile blobs to the
+    /// remote-execution endpoint.
+    /// \returns The hash of the resulting tree if all distfile blobs found,
+    /// nullopt otherwise.
+    [[nodiscard]] auto ServeDistdirTree(
+        std::shared_ptr<std::unordered_map<std::string, std::string>> const&
+            distfiles,
+        bool sync_tree) -> std::optional<std::string>;
+
     /// \brief Make a given content blob available in remote CAS, if known by
     /// serve remote.
     /// \param[in] content Hash of the archive content to look up.