fix: index branch name as well & fix incorrect file url when not main

zchee · zchee · commit ceb858b611c1 · 2025-09-11T09:04:33.000+09:00
Signed-off-by: Koichi Shiraishi &lt;zchee.io@gmail.com&gt;
diff --git a/packages/chrome-extension/src/background.ts b/packages/chrome-extension/src/background.ts
@@ -124,6 +124,7 @@ class MilvusVectorDB {
             return results.map(result => ({
                 id: result.id,
                 content: result.content,
+                branch: result.branch, // Include branch name for frontend display
                 relativePath: result.relativePath,
                 startLine: result.startLine,
                 endLine: result.endLine,
@@ -310,6 +311,27 @@ async function handleRateLimit(response: Response): Promise<void> {
     }
 }
 
+async function fetchRepoBranch(owner: string, repo: string): Promise<string> {
+    const token = await getGitHubToken();
+
+    // First get the default branch
+    const repoInfoUrl = `https://api.github.com/repos/${owner}/${repo}`;
+    const repoResponse = await fetch(repoInfoUrl, {
+        headers: {
+            'Authorization': `Bearer ${token}`,
+            'Accept': 'application/vnd.github+json',
+            'X-GitHub-Api-Version': '2022-11-28'
+        }
+    });
+
+    if (!repoResponse.ok) {
+        throw new Error(`GitHub API error: ${repoResponse.status} - ${await repoResponse.text()}`);
+    }
+
+    const repoData = await repoResponse.json();
+    return repoData.default_branch || 'main';
+}
+
 async function fetchRepoFiles(owner: string, repo: string): Promise<any[]> {
     const token = await getGitHubToken();
 
@@ -457,11 +479,14 @@ async function handleIndexRepo(request: any, sendResponse: Function) {
         const chunkOverlap = 200;  // Same as VSCode extension default
 
         // Fetch repository files
-        const files = await fetchRepoFiles(owner, repo);
+        const [branch, files] = await Promise.all([
+            fetchRepoBranch(owner, repo),
+            fetchRepoFiles(owner, repo),
+        ]);
         console.log(`Found ${files.length} files to index`);
 
         // Process files using core package approach
-        const result = await processFileList(files, owner, repo, repoId, vectorDB, chunkSize, chunkOverlap);
+        const result = await processFileList(files, owner, repo, branch, repoId, vectorDB, chunkSize, chunkOverlap);
 
         // Send completion message
         chrome.tabs.query({ active: true, currentWindow: true }, (tabs) => {
@@ -499,6 +524,7 @@ async function processFileList(
     files: any[],
     owner: string,
     repo: string,
+    branch: string,
     repoId: string,
     vectorDB: MilvusVectorDB,
     chunkSize: number,
@@ -532,6 +558,7 @@ async function processFileList(
                     const codeChunk: CodeChunk = {
                         id: `${file.path}_chunk_${j}`,
                         content: chunk.content,
+                        branch: branch,
                         relativePath: file.path,
                         startLine: chunk.startLine,
                         endLine: chunk.endLine,
diff --git a/packages/chrome-extension/src/content.ts b/packages/chrome-extension/src/content.ts
@@ -366,6 +366,7 @@ function displayResults(results: any[]) {
         const [owner, repo] = window.location.pathname.slice(1).split('/');
 
         // Format the file path to show it nicely
+        const branch = result.branch;
         const filePath = result.relativePath;
         const fileExt = filePath.split('.').pop();
 
@@ -383,7 +384,7 @@ function displayResults(results: any[]) {
                 <svg class="octicon mr-2 color-fg-muted" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16">
                     <path fill-rule="evenodd" d="M3.75 1.5a.25.25 0 00-.25.25v11.5c0 .138.112.25.25.25h8.5a.25.25 0 00.25-.25V6H9.75A1.75 1.75 0 018 4.25V1.5H3.75zm5.75.56v2.19c0 .138.112.25.25.25h2.19L9.5 2.06zM2 1.75C2 .784 2.784 0 3.75 0h5.086c.464 0 .909.184 1.237.513l3.414 3.414c.329.328.513.773.513 1.237v8.086A1.75 1.75 0 0112.25 15h-8.5A1.75 1.75 0 012 13.25V1.75z"></path>
                 </svg>
-                <a href="https://github.com/${owner}/${repo}/blob/main/${result.relativePath}#L${result.startLine}" class="Link--primary flex-auto" style="font-weight: 600;">
+                <a href="https://github.com/${owner}/${repo}/blob/${branch}/${result.relativePath}#L${result.startLine}" class="Link--primary flex-auto" style="font-weight: 600;">
                     ${result.relativePath}
                 </a>
                 <span class="Label Label--secondary ml-1">${fileExt}</span>
@@ -603,4 +604,4 @@ new MutationObserver((mutations, observer) => {
         // Just check if UI needs to be injected (for dynamic content)
         injectUI();
     }
-}).observe(document.body, { childList: true, subtree: true }); 
+}).observe(document.body, { childList: true, subtree: true }); 
diff --git a/packages/chrome-extension/src/milvus/chromeMilvusAdapter.ts b/packages/chrome-extension/src/milvus/chromeMilvusAdapter.ts
@@ -9,6 +9,7 @@ import { MilvusRestfulVectorDatabase } from '../stubs/milvus-vectordb-stub';
 export interface CodeChunk {
     id: string;
     content: string;
+    branch: string;
     relativePath: string;
     startLine: number;
     endLine: number;
@@ -20,6 +21,7 @@ export interface CodeChunk {
 export interface SearchResult {
     id: string;
     content: string;
+    branch: string;
     relativePath: string;
     startLine: number;
     endLine: number;
@@ -113,6 +115,7 @@ export class ChromeMilvusAdapter {
             id: chunk.id,
             vector: chunk.vector || [],
             content: chunk.content,
+            branch: chunk.branch,
             relativePath: chunk.relativePath,
             startLine: chunk.startLine,
             endLine: chunk.endLine,
@@ -149,6 +152,7 @@ export class ChromeMilvusAdapter {
             const searchResults = results.map(result => ({
                 id: result.document.id,
                 content: result.document.content,
+                branch: result.document.branch,
                 relativePath: result.document.relativePath,
                 startLine: result.document.startLine,
                 endLine: result.document.endLine,
diff --git a/packages/chrome-extension/src/stubs/milvus-vectordb-stub.ts b/packages/chrome-extension/src/stubs/milvus-vectordb-stub.ts
@@ -6,6 +6,7 @@ export interface VectorDocument {
     id: string;
     vector: number[];
     content: string;
+    branch: string;
     relativePath: string;
     startLine: number;
     endLine: number;
@@ -141,6 +142,13 @@ export class MilvusRestfulVectorDatabase {
                                 max_length: 65535
                             }
                         },
+                        {
+                            fieldName: "branch",
+                            dataType: "VarChar",
+                            elementTypeParams: {
+                                max_length: 512
+                            }
+                        },
                         {
                             fieldName: "relativePath",
                             dataType: "VarChar",
@@ -247,6 +255,7 @@ export class MilvusRestfulVectorDatabase {
                 id: doc.id,
                 vector: doc.vector,
                 content: doc.content,
+                branch: doc.branch,
                 relativePath: doc.relativePath,
                 startLine: doc.startLine,
                 endLine: doc.endLine,
@@ -280,6 +289,7 @@ export class MilvusRestfulVectorDatabase {
                 limit: topK,
                 outputFields: [
                     "content",
+                    "branch",
                     "relativePath",
                     "startLine",
                     "endLine",
@@ -308,6 +318,7 @@ export class MilvusRestfulVectorDatabase {
                         id: item.id?.toString() || '',
                         vector: queryVector,
                         content: item.content || '',
+                        branch: item.branch || '',
                         relativePath: item.relativePath || '',
                         startLine: item.startLine || 0,
                         endLine: item.endLine || 0,
