feat(mention): allow selection of context from branches for remote files and directories

lib/prompt-editor/src/mentions/mentionMenu/MentionMenu.branch.test.ts

@@ -0,0 +1,161 @@
+import { ContextItemSource, REMOTE_DIRECTORY_PROVIDER_URI } from '@sourcegraph/cody-shared'
+import type { MentionMenuData } from '@sourcegraph/cody-shared'
+import { describe, expect, test } from 'vitest'
+import { URI } from 'vscode-uri'
+
+// This would be imported from the MentionMenu component if it were exported
+// For now, we'll create a test version of the function
+function getBranchHelpText(
+    items: NonNullable<MentionMenuData['items']>,
+    mentionQuery: { text: string }
+): string {
+    // Check if we have branch information from the current search
+    const firstItem = items[0]
+    if (firstItem?.type === 'openctx') {
+        const openCtxItem = firstItem as any // Simplified for testing
+        if (openCtxItem.mention?.data?.branch) {
+            return `* Sourced from the '${openCtxItem.mention.data.branch}' branch`
+        }
+    }
+
+    // Check if user has specified a branch in the query
+    if (mentionQuery.text.includes('@')) {
+        const branchPart = mentionQuery.text.split('@')[1]
+        if (branchPart) {
+            // Remove anything after colon (directory path)
+            const branchName = branchPart.split(':')[0]
+            return `* Sourced from the '${branchName}' branch`
+        }
+    }
+
+    return '* Sourced from the remote default branch'
+}
+
+describe('MentionMenu branch selection', () => {
+    test('should show default branch text when no branch is specified', () => {
+        const items: MentionMenuData['items'] = [
+            {
+                type: 'openctx',
+                provider: 'openctx',
+                title: 'src/components',
+                uri: URI.parse('https://example.com/repo/-/tree/src/components'),
+                providerUri: REMOTE_DIRECTORY_PROVIDER_URI,
+                source: ContextItemSource.User,
+                mention: {
+                    uri: 'https://example.com/repo/-/tree/src/components',
+                    data: {
+                        repoName: 'test-repo',
+                        directoryPath: 'src/components',
+                    },
+                },
+            },
+        ]
+
+        const mentionQuery = { text: 'test-repo:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe('* Sourced from the remote default branch')
+    })
+
+    test('should show branch name when branch is specified in query', () => {
+        const items: MentionMenuData['items'] = [
+            {
+                type: 'openctx',
+                provider: 'openctx',
+                title: 'src/components',
+                uri: URI.parse('https://example.com/repo/-/tree/src/components'),
+                providerUri: REMOTE_DIRECTORY_PROVIDER_URI,
+                source: ContextItemSource.User,
+                mention: {
+                    uri: 'https://example.com/repo/-/tree/src/components',
+                    data: {
+                        repoName: 'test-repo',
+                        directoryPath: 'src/components',
+                    },
+                },
+            },
+        ]
+
+        const mentionQuery = { text: 'test-repo@feature-branch:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe("* Sourced from the 'feature-branch' branch")
+    })
+
+    test('should show branch name from mention data when available', () => {
+        const items: MentionMenuData['items'] = [
+            {
+                type: 'openctx',
+                provider: 'openctx',
+                title: 'src/components',
+                uri: URI.parse('https://example.com/repo/-/tree/src/components'),
+                providerUri: REMOTE_DIRECTORY_PROVIDER_URI,
+                source: ContextItemSource.User,
+                mention: {
+                    uri: 'https://example.com/repo/-/tree/src/components',
+                    data: {
+                        repoName: 'test-repo',
+                        directoryPath: 'src/components',
+                        branch: 'main',
+                    },
+                },
+            },
+        ]
+
+        const mentionQuery = { text: 'test-repo:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe("* Sourced from the 'main' branch")
+    })
+
+    test('should prefer mention data branch over query branch', () => {
+        const items: MentionMenuData['items'] = [
+            {
+                type: 'openctx',
+                provider: 'openctx',
+                title: 'src/components',
+                uri: URI.parse('https://example.com/repo/-/tree/src/components'),
+                providerUri: REMOTE_DIRECTORY_PROVIDER_URI,
+                source: ContextItemSource.User,
+                mention: {
+                    uri: 'https://example.com/repo/-/tree/src/components',
+                    data: {
+                        repoName: 'test-repo',
+                        directoryPath: 'src/components',
+                        branch: 'actual-branch',
+                    },
+                },
+            },
+        ]
+
+        const mentionQuery = { text: 'test-repo@query-branch:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe("* Sourced from the 'actual-branch' branch")
+    })
+
+    test('should handle empty items array', () => {
+        const items: MentionMenuData['items'] = []
+        const mentionQuery = { text: 'test-repo@main:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe("* Sourced from the 'main' branch")
+    })
+
+    test('should handle non-openctx items', () => {
+        const items: MentionMenuData['items'] = [
+            {
+                type: 'file',
+                provider: 'file',
+                title: 'test.ts',
+                uri: URI.parse('file:///test.ts'),
+                source: ContextItemSource.User,
+            },
+        ]
+
+        const mentionQuery = { text: 'test-repo@feature:src' }
+
+        const result = getBranchHelpText(items!, mentionQuery)
+        expect(result).toBe("* Sourced from the 'feature' branch")
+    })
+})

lib/prompt-editor/src/mentions/mentionMenu/MentionMenu.tsx

@@ -305,7 +305,7 @@ export const MentionMenu: FunctionComponent<
                                 'tw-bg-accent'
                             )}
                         >
-                            * Sourced from the remote default branch
+                            {getBranchHelpText(data.items, mentionQuery)}
                         </CommandLoading>
                     )}
 
@@ -341,6 +341,35 @@ function commandRowValue(
     return contextItemID(row)
 }
 
+function getBranchHelpText(
+    items: NonNullable<MentionMenuData['items']>,
+    mentionQuery: MentionQuery
+): string {
+    // Check if we have branch information from the current search
+    const firstItem = items[0]
+    if (firstItem?.type === 'openctx') {
+        const openCtxItem = firstItem as ContextItem & {
+            type: 'openctx'
+            mention?: { data?: { branch?: string } }
+        }
+        if (openCtxItem.mention?.data?.branch) {
+            return `* Sourced from the '${openCtxItem.mention.data.branch}' branch`
+        }
+    }
+
+    // Check if user has specified a branch in the query
+    if (mentionQuery.text.includes('@')) {
+        const branchPart = mentionQuery.text.split('@')[1]
+        if (branchPart) {
+            // Remove anything after colon (directory path)
+            const branchName = branchPart.split(':')[0]
+            return `* Sourced from the '${branchName}' branch`
+        }
+    }
+
+    return '* Sourced from the remote default branch'
+}
+
 function getEmptyLabel(
     parentItem: ContextMentionProviderMetadata | null,
     mentionQuery: MentionQuery

lib/shared/src/sourcegraph-api/graphql/client.ts

@@ -42,6 +42,7 @@ import {
     CURRENT_USER_INFO_QUERY,
     CURRENT_USER_ROLE_QUERY,
     DELETE_ACCESS_TOKEN_MUTATION,
+    DIRECTORY_CONTENTS_QUERY,
     EVALUATE_FEATURE_FLAGS_QUERY,
     EVALUATE_FEATURE_FLAG_QUERY,
     FILE_CONTENTS_QUERY,
@@ -370,6 +371,23 @@ interface FileContentsResponse {
     } | null
 }
 
+interface DirectoryContentsResponse {
+    repository: {
+        commit: {
+            tree: {
+                entries: Array<{
+                    name: string
+                    path: string
+                    byteSize?: number
+                    url: string
+                    content?: string
+                    isDirectory?: boolean
+                }>
+            } | null
+        } | null
+    } | null
+}
+
 export interface RepositoryIdResponse {
     repository: { id: string } | null
 }
@@ -1049,6 +1067,21 @@ export class SourcegraphGraphQLAPIClient {
         }).then(response => extractDataOrError(response, data => data))
     }
 
+    public async getDirectoryContents(
+        repoName: string,
+        path: string,
+        revision = 'HEAD'
+    ): Promise<DirectoryContentsResponse | Error> {
+        return this.fetchSourcegraphAPI<APIResponse<DirectoryContentsResponse>>(
+            DIRECTORY_CONTENTS_QUERY,
+            {
+                repoName,
+                path,
+                revision,
+            }
+        ).then(response => extractDataOrError(response, data => data))
+    }
+
     public async getRepoId(repoName: string): Promise<string | null | Error> {
         return this.fetchSourcegraphAPI<APIResponse<RepositoryIdResponse>>(REPOSITORY_ID_QUERY, {
             name: repoName,

lib/shared/src/sourcegraph-api/graphql/queries.ts

@@ -754,3 +754,27 @@ export const NLS_SEARCH_QUERY = `
             }
         }
     }`
+
+export const DIRECTORY_CONTENTS_QUERY = `
+query GetDirectoryContents($repoName: String!, $revision: String!, $path: String!) {
+  repository(name: $repoName) {
+    commit(rev: $revision) {
+      tree(path: $path) {
+        entries(first: 50) {
+          ... on GitBlob {
+            name
+            path
+            byteSize
+            url
+            content
+          }
+          ... on GitTree {
+            name
+            path
+            isDirectory
+          }
+        }
+      }
+    }
+  }
+}`