From 237345bf904d2d2afcdcf5ca950861350d6bccc6 Mon Sep 17 00:00:00 2001
From: Weixuan Fu <weixuanf@outlook.com>
Date: Sun, 31 Mar 2024 19:27:52 +0800
Subject: [PATCH] return nodeType to mapping instead of set

---
 service/node_service.py                     |  9 ++++-----
 ui/src/spacejson/handleDownloadSpaceJson.ts | 22 +++++++++++++++------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/service/node_service.py b/service/node_service.py
index e34c5a6..7ad17dd 100644
--- a/service/node_service.py
+++ b/service/node_service.py
@@ -172,7 +172,7 @@ def get_git_repo(node_type: str):
         git_repo = git_repo[:-4]
     username = git_repo.split("/")[-2]
     repo_name = git_repo.split("/")[-1]
-    return {"repoID": f"{username}/{repo_name}", "commitHash": commit_hash}
+    return {"gitRepo": f"{username}/{repo_name}", "commitHash": commit_hash}
 
 @server.PromptServer.instance.routes.post("/workspace/fetch_node_repos")  # Handle POST requests
 async def fetch_node_repos(request):
@@ -180,14 +180,13 @@ async def fetch_node_repos(request):
     nodetypes = data.get("nodes")
     if not nodetypes:
         return web.Response(status=400, text="NodeTypes parameter is required and should be a list of node types.")
-    repos = {}
+    repos_mapping = {}
     for nodetype in nodetypes:
         try:
             repo = get_git_repo(nodetype)
             if repo:
-                repos[repo.get('repoID')] = repo
+                repos_mapping[nodetype] = repo
         except Exception as e:
             print(f"Error fetching repo for {nodetype}: {e}")
-    result = list(repos.values())
-    return web.Response(text=json.dumps(result), content_type='application/json')
+    return web.Response(text=json.dumps(repos_mapping), content_type='application/json')
 
diff --git a/ui/src/spacejson/handleDownloadSpaceJson.ts b/ui/src/spacejson/handleDownloadSpaceJson.ts
index a1cf314..8b146da 100644
--- a/ui/src/spacejson/handleDownloadSpaceJson.ts
+++ b/ui/src/spacejson/handleDownloadSpaceJson.ts
@@ -32,7 +32,7 @@ type ImageFile = {
 };
 type NodeRepo = {
   commitHash: string;
-  repoID: string;
+  gitRepo: string;
 };
 
 export type DepsResult = {
@@ -114,10 +114,13 @@ export async function extractAndFetchFileNames(
       }
     },
   );
-  const resp = await fetch("workspace/fetch_node_repos", {
-    method: "POST",
-    body: JSON.stringify({ nodes: nodes.map((n) => n.type) }),
-  })
+  const reposMapping: Record<string, NodeRepo> = await fetch(
+    "workspace/fetch_node_repos",
+    {
+      method: "POST",
+      body: JSON.stringify({ nodes: nodes.map((n) => n.type) }),
+    },
+  )
     .then((res) => {
       if (!res.ok) {
         return [];
@@ -132,7 +135,14 @@ export async function extractAndFetchFileNames(
   const models = await Promise.all(modelPromises);
   const modelsMap: Record<string, DepsResultModel> = {};
   models.forEach((model) => {
+    // if (model.nodeType) {
+    //   model.gitRepo = reposMapping[model.nodeType];
+    // }
     modelsMap[model.filename] = model;
   });
-  return { models: modelsMap, images, nodeRepos: resp };
+  const nodeRepos: Record<string, NodeRepo> = {};
+  Object.values(reposMapping).forEach((repo) => {
+    nodeRepos[repo.gitRepo] = repo;
+  });
+  return { models: modelsMap, images, nodeRepos: Object.values(nodeRepos) };
 }