From b4eb6ed1dda6f2ffd4a014c5ff2f26abcc56592a Mon Sep 17 00:00:00 2001
From: Naman Kumar <naman@outlook.in>
Date: Fri, 27 Dec 2024 15:23:23 +0530
Subject: [PATCH] Open remote files locally in VSCode

---
 vscode/src/chat/chat-view/ChatController.ts   | 13 +++++
 .../chat/chat-view/sourcegraphRemoteFile.ts   | 57 +++++++++++++++++++
 vscode/src/chat/protocol.ts                   |  1 +
 vscode/src/main.ts                            | 11 +++-
 .../components/codeSnippet/CodeSnippet.tsx    | 20 +++++++
 .../components/FileMatchChildren.tsx          | 16 +++++-
 .../codeSnippet/components/RepoLink.tsx       | 11 +++-
 7 files changed, 125 insertions(+), 4 deletions(-)
 create mode 100644 vscode/src/chat/chat-view/sourcegraphRemoteFile.ts

diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts
index de853295ecbd..6931f677834d 100644
--- a/vscode/src/chat/chat-view/ChatController.ts
+++ b/vscode/src/chat/chat-view/ChatController.ts
@@ -361,6 +361,9 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                     viewColumn: vscode.ViewColumn.Beside,
                 })
                 break
+            case 'openRemoteFile':
+                this.openRemoteFile(message.uri)
+                break
             case 'newFile':
                 await handleCodeFromSaveToNewFile(message.text, this.editor)
                 break
@@ -910,6 +913,16 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         return
     }
 
+    private async openRemoteFile(uri: vscode.Uri) {
+        const json = uri.toJSON()
+        json.scheme = 'sourcegraph'
+        const sourcegraphSchemaURI = vscode.Uri.from(json)
+
+        vscode.workspace
+            .openTextDocument(sourcegraphSchemaURI)
+            .then(doc => vscode.window.showTextDocument(doc))
+    }
+
     private submitOrEditOperation: AbortController | undefined
     public startNewSubmitOrEditOperation(): AbortSignal {
         this.submitOrEditOperation?.abort()
diff --git a/vscode/src/chat/chat-view/sourcegraphRemoteFile.ts b/vscode/src/chat/chat-view/sourcegraphRemoteFile.ts
new file mode 100644
index 000000000000..375dbe207184
--- /dev/null
+++ b/vscode/src/chat/chat-view/sourcegraphRemoteFile.ts
@@ -0,0 +1,57 @@
+import { graphqlClient, isError } from '@sourcegraph/cody-shared'
+import { LRUCache } from 'lru-cache'
+import * as vscode from 'vscode'
+
+export class SourcegraphRemoteFileProvider
+    implements vscode.TextDocumentContentProvider, vscode.Disposable
+{
+    private cache = new LRUCache<string, string>({ max: 128 })
+    private disposables: vscode.Disposable[] = []
+
+    constructor() {
+        this.disposables.push(vscode.workspace.registerTextDocumentContentProvider('sourcegraph', this))
+    }
+
+    async provideTextDocumentContent(uri: vscode.Uri): Promise<string> {
+        const content =
+            this.cache.get(uri.toString()) ||
+            (await SourcegraphRemoteFileProvider.getFileContentsFromURL(uri))
+
+        this.cache.set(uri.toString(), content)
+
+        return content
+    }
+
+    public dispose(): void {
+        this.cache.clear()
+        for (const disposable of this.disposables) {
+            disposable.dispose()
+        }
+        this.disposables = []
+    }
+
+    private static async getFileContentsFromURL(URL: vscode.Uri): Promise<string> {
+        const path = URL.path
+        const [repoRev = '', filePath] = path.split('/-/blob/')
+        let [repoName, rev = 'HEAD'] = repoRev.split('@')
+        repoName = repoName.replace(/^\/+/, '')
+
+        if (!repoName || !filePath) {
+            throw new Error('Invalid URI')
+        }
+
+        const dataOrError = await graphqlClient.getFileContents(repoName, filePath, rev)
+
+        if (isError(dataOrError)) {
+            throw new Error(dataOrError.message)
+        }
+
+        const content = dataOrError.repository?.commit?.file?.content
+
+        if (!content) {
+            throw new Error('File not found')
+        }
+
+        return content
+    }
+}
diff --git a/vscode/src/chat/protocol.ts b/vscode/src/chat/protocol.ts
index f7004c609cb2..57274b8302b5 100644
--- a/vscode/src/chat/protocol.ts
+++ b/vscode/src/chat/protocol.ts
@@ -74,6 +74,7 @@ export type WebviewMessage =
     | { command: 'restoreHistory'; chatID: string }
     | { command: 'links'; value: string }
     | { command: 'openURI'; uri: Uri }
+    | { command: 'openRemoteFile'; uri: Uri }
     | {
           command: 'openFileLink'
           uri: Uri
diff --git a/vscode/src/main.ts b/vscode/src/main.ts
index 5ba73e55b4e5..b694c3882758 100644
--- a/vscode/src/main.ts
+++ b/vscode/src/main.ts
@@ -51,6 +51,7 @@ import { registerAutoEditTestRenderCommand } from './autoedits/renderer/mock-ren
 import type { MessageProviderOptions } from './chat/MessageProvider'
 import { ChatsController, CodyChatEditorViewType } from './chat/chat-view/ChatsController'
 import { ContextRetriever } from './chat/chat-view/ContextRetriever'
+import { SourcegraphRemoteFileProvider } from './chat/chat-view/sourcegraphRemoteFile'
 import type { ChatIntentAPIClient } from './chat/context/chatIntentAPIClient'
 import {
     ACCOUNT_LIMITS_INFO_URL,
@@ -905,7 +906,15 @@ function registerChat(
         extensionClient: platform.extensionClient,
     })
     const promptsManager = new PromptsManager({ chatsController })
-    disposables.push(ghostHintDecorator, editorManager, new CodeActionProvider(), promptsManager)
+    const sourcegraphRemoteFileProvider = new SourcegraphRemoteFileProvider()
+
+    disposables.push(
+        ghostHintDecorator,
+        editorManager,
+        new CodeActionProvider(),
+        promptsManager,
+        sourcegraphRemoteFileProvider
+    )
 
     // Register a serializer for reviving the chat panel on reload
     if (vscode.window.registerWebviewPanelSerializer) {
diff --git a/vscode/webviews/components/codeSnippet/CodeSnippet.tsx b/vscode/webviews/components/codeSnippet/CodeSnippet.tsx
index 9cd8b0674ba7..4eca110fae87 100644
--- a/vscode/webviews/components/codeSnippet/CodeSnippet.tsx
+++ b/vscode/webviews/components/codeSnippet/CodeSnippet.tsx
@@ -24,12 +24,16 @@ import {
     pluralize,
 } from './utils'
 
+import { CodyIDE } from '@sourcegraph/cody-shared'
 import type {
     NLSSearchFileMatch,
     NLSSearchResult,
 } from '@sourcegraph/cody-shared/src/sourcegraph-api/graphql/client'
 import type { Observable } from 'observable-fns'
 import { useInView } from 'react-intersection-observer'
+import { URI } from 'vscode-uri'
+import { getVSCodeAPI } from '../../utils/VSCodeApi'
+import { useConfig } from '../../utils/useConfig'
 import styles from './CodeSnippet.module.css'
 
 const DEFAULT_VISIBILITY_OFFSET = '500px'
@@ -123,6 +127,20 @@ export const FileMatchSearchResult: FC<PropsWithChildren<FileMatchSearchResultPr
     const expandable = !showAllMatches && expandedHighlightCount > collapsedHighlightCount
 
     useEffect(() => setExpanded(allExpanded || defaultExpanded), [allExpanded, defaultExpanded])
+    const {
+        clientCapabilities: { agentIDE },
+    } = useConfig()
+    const openRemoteFile = useCallback(() => {
+        if (agentIDE !== CodyIDE.VSCode) {
+            return
+        }
+
+        const uri = URI.parse(fileURL)
+        getVSCodeAPI().postMessage({
+            command: 'openRemoteFile',
+            uri,
+        })
+    }, [fileURL, agentIDE])
 
     const handleVisibility = useCallback(
         (inView: boolean, entry: IntersectionObserverEntry) => {
@@ -183,6 +201,7 @@ export const FileMatchSearchResult: FC<PropsWithChildren<FileMatchSearchResultPr
             repoName={result.repository.name}
             repoURL={repoAtRevisionURL}
             filePath={result.file.path}
+            onFilePathClick={openRemoteFile}
             pathMatchRanges={result.pathMatches ?? []}
             fileURL={fileURL}
             repoDisplayName={
@@ -231,6 +250,7 @@ export const FileMatchSearchResult: FC<PropsWithChildren<FileMatchSearchResultPr
                     serverEndpoint={serverEndpoint}
                     result={result}
                     grouped={expanded ? expandedGroups : collapsedGroups}
+                    onLineClick={openRemoteFile}
                 />
                 {expandable && (
                     <button
diff --git a/vscode/webviews/components/codeSnippet/components/FileMatchChildren.tsx b/vscode/webviews/components/codeSnippet/components/FileMatchChildren.tsx
index 67415f4f2f9b..2814f04d960b 100644
--- a/vscode/webviews/components/codeSnippet/components/FileMatchChildren.tsx
+++ b/vscode/webviews/components/codeSnippet/components/FileMatchChildren.tsx
@@ -8,8 +8,10 @@ import { getFileMatchUrl } from '../utils'
 
 import { CodeExcerpt } from './CodeExcerpt'
 
+import { CodyIDE } from '@sourcegraph/cody-shared'
 import type { NLSSearchFileMatch } from '@sourcegraph/cody-shared/src/sourcegraph-api/graphql/client'
 import { getVSCodeAPI } from '../../../utils/VSCodeApi'
+import { useConfig } from '../../../utils/useConfig'
 import resultStyles from '../CodeSnippet.module.css'
 import styles from './FileMatchChildren.module.css'
 
@@ -17,10 +19,11 @@ interface FileMatchProps {
     result: NLSSearchFileMatch
     grouped: MatchGroup[]
     serverEndpoint: string
+    onLineClick?: () => void
 }
 
 export const FileMatchChildren: FC<PropsWithChildren<FileMatchProps>> = props => {
-    const { result, grouped, serverEndpoint } = props
+    const { result, grouped, serverEndpoint, onLineClick } = props
 
     const createCodeExcerptLink = (group: MatchGroup): string => {
         const urlBuilder = SourcegraphURL.from(getFileMatchUrl(serverEndpoint, result))
@@ -39,15 +42,24 @@ export const FileMatchChildren: FC<PropsWithChildren<FileMatchProps>> = props =>
         return urlBuilder.toString()
     }
 
+    const {
+        clientCapabilities: { agentIDE },
+    } = useConfig()
+
     const navigateToFile = useCallback(
         (line: number) => {
+            if (agentIDE === CodyIDE.VSCode && onLineClick) {
+                onLineClick()
+                return
+            }
+
             // TODO: this does not work on web as opening links from within a web worker does not work.
             getVSCodeAPI().postMessage({
                 command: 'links',
                 value: `${getFileMatchUrl(serverEndpoint, result)}?L${line}`,
             })
         },
-        [serverEndpoint, result]
+        [serverEndpoint, result, onLineClick, agentIDE]
     )
 
     return (
diff --git a/vscode/webviews/components/codeSnippet/components/RepoLink.tsx b/vscode/webviews/components/codeSnippet/components/RepoLink.tsx
index 5b1cb2560f87..9eb2576272d7 100644
--- a/vscode/webviews/components/codeSnippet/components/RepoLink.tsx
+++ b/vscode/webviews/components/codeSnippet/components/RepoLink.tsx
@@ -1,7 +1,9 @@
 import type * as React from 'react'
 import { useEffect, useRef } from 'react'
 
+import { CodyIDE } from '@sourcegraph/cody-shared'
 import { ChevronDown, ChevronUp } from 'lucide-react'
+import { useConfig } from '../../../utils/useConfig'
 import { cn } from '../../shadcn/utils'
 import { highlightNode } from '../highlights'
 import type { Range } from '../types'
@@ -17,6 +19,7 @@ interface Props {
     isKeyboardSelectable?: boolean
     collapsed: boolean
     onToggleCollapse: () => void
+    onFilePathClick?: () => void
 }
 
 /**
@@ -35,6 +38,7 @@ export const RepoFileLink: React.FunctionComponent<React.PropsWithChildren<Props
         isKeyboardSelectable,
         collapsed,
         onToggleCollapse,
+        onFilePathClick,
     } = props
 
     const [fileBase, fileName] = splitPath(filePath)
@@ -52,6 +56,10 @@ export const RepoFileLink: React.FunctionComponent<React.PropsWithChildren<Props
         }
     }, [pathMatchRanges, fileName])
 
+    const {
+        clientCapabilities: { agentIDE },
+    } = useConfig()
+
     return (
         <span className={cn(className, 'tw-flex tw-items-center tw-w-full')}>
             {collapsed ? (
@@ -73,11 +81,12 @@ export const RepoFileLink: React.FunctionComponent<React.PropsWithChildren<Props
                 </a>
                 <span aria-hidden={true}> ›</span>{' '}
                 <a
-                    href={fileURL}
+                    href={agentIDE === CodyIDE.VSCode ? '' : fileURL}
                     ref={containerElement}
                     target="_blank"
                     rel="noreferrer"
                     data-selectable-search-result={isKeyboardSelectable}
+                    onClick={onFilePathClick}
                 >
                     {fileBase ? `${fileBase}/` : null}
                     <strong>{fileName}</strong>