added(webviews): Adds protocol based webview loading (CODY-3536)

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ClientCapabilities.kt

@@ -20,7 +20,8 @@ data class ClientCapabilities(
   val globalState: GlobalStateEnum? = null, // Oneof: stateless, server-managed, client-managed
   val secrets: SecretsEnum? = null, // Oneof: stateless, client-managed
   val webview: WebviewEnum? = null, // Oneof: agentic, native
-  val webviewNativeConfig: WebviewNativeConfigParams? = null,
+  val uriSchemeLoaders: List<String>? = null,
+  val webviewNativeConfig: WebviewNativeConfig? = null,
 ) {
 
   enum class CompletionsEnum {

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt

@@ -22,6 +22,8 @@ interface CodyAgentClient {
   fun textDocument_show(params: TextDocument_ShowParams): CompletableFuture<Boolean>
   @JsonRequest("workspace/edit")
   fun workspace_edit(params: WorkspaceEditParams): CompletableFuture<Boolean>
+  @JsonRequest("uri/readUTF8")
+  fun uri_readUTF8(params: Uri_ReadUTF8Params): CompletableFuture<Uri_ReadUTF8Result>
   @JsonRequest("secrets/get")
   fun secrets_get(params: Secrets_GetParams): CompletableFuture<String?>
   @JsonRequest("secrets/store")

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Constants.kt

@@ -41,6 +41,7 @@ object Constants {
   const val fetching = "fetching"
   const val file = "file"
   const val free = "free"
+  const val fs = "fs"
   const val function = "function"
   const val gateway = "gateway"
   const val history = "history"
@@ -89,5 +90,6 @@ object Constants {
   const val use = "use"
   const val user = "user"
   const val warning = "warning"
+  const val webviewasset = "webviewasset"
   const val workspace = "workspace"
 }

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Uri_ReadUTF8Params.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.agent.protocol_generated;
+
+data class Uri_ReadUTF8Params(
+  val uri: String,
+)
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Uri_ReadUTF8Result.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.agent.protocol_generated;
+
+data class Uri_ReadUTF8Result(
+  val text: String,
+)
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/WebviewNativeConfigParams.kt → agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/WebviewNativeConfig.kt

@@ -3,9 +3,10 @@ package com.sourcegraph.cody.agent.protocol_generated;
 
 import com.google.gson.annotations.SerializedName;
 
-data class WebviewNativeConfigParams(
+data class WebviewNativeConfig(
   val view: ViewEnum, // Oneof: multiple, single
   val cspSource: String,
+  val assetLoader: AssetLoaderEnum? = null, // Oneof: fs, webviewasset
   val webviewBundleServingPrefix: String,
   val rootDir: String? = null,
   val injectScript: String? = null,
@@ -16,5 +17,10 @@ data class WebviewNativeConfigParams(
     @SerializedName("multiple") Multiple,
     @SerializedName("single") Single,
   }
+
+  enum class AssetLoaderEnum {
+    @SerializedName("fs") Fs,
+    @SerializedName("webviewasset") Webviewasset,
+  }
 }
 

agent/src/NativeWebview.ts

@@ -1,18 +1,16 @@
 import * as uuid from 'uuid'
 import * as vscode from 'vscode'
 import type { Agent } from './agent'
-import type { DefiniteWebviewOptions } from './protocol-alias'
+import type { DefiniteWebviewOptions, WebviewNativeConfig } from './protocol-alias'
 import * as vscode_shim from './vscode-shim'
 
-export type NativeWebviewConfig = { cspSource: string; webviewBundleServingPrefix: string }
-
 type NativeWebviewHandle = string
 
 /**
  * A delegate for adapting the VSCode Webview, WebviewPanel and WebviewView API
  * to a client which has a native webview implementation.
  */
-interface WebviewProtocolDelegate {
+interface WebviewProtocolDelegate extends WebviewNativeConfig {
     // CSP, resource-related
     readonly webviewBundleLocalPrefix: vscode.Uri
     readonly webviewBundleServingPrefix: string
@@ -88,15 +86,14 @@ export function resolveWebviewView(
 export function registerNativeWebviewHandlers(
     agent: Agent,
     webviewBundleLocalPrefix: vscode.Uri,
-    config: NativeWebviewConfig
+    config: WebviewNativeConfig
 ): void {
     webviewProtocolDelegate = {
+        ...config,
         // TODO: When we want to serve resources outside dist/, make Agent
         // include 'dist' in its bundle paths, and simply set this to
         // extensionUri.
         webviewBundleLocalPrefix,
-        webviewBundleServingPrefix: config.webviewBundleServingPrefix,
-        cspSource: config.cspSource,
         createWebviewPanel: (handle, viewType, title, showOptions, options) => {
             agent.notify('webview/createWebviewPanel', {
                 handle,

agent/src/agent.ts

@@ -1478,6 +1478,21 @@ export class Agent extends MessageHandler implements ExtensionClient {
         return result ? vscode_shim.workspace.openTextDocument(result.uri) : undefined
     }
 
+    public async readUriUTF8(uri: vscode.Uri): Promise<string | undefined> {
+        if (this.clientInfo?.capabilities?.uriSchemeLoaders?.includes(uri.scheme)) {
+            const { text } = await this.request('uri/readUTF8', {
+                uri: uri.toString(),
+            })
+
+            return text
+        }
+        const errorMessage =
+            `Client does not support ${uri.scheme} documents. To fix this problem, add ${uri.scheme} to ` +
+            `ClientCapabilities.uriSchemeLoaders and implement the ${uri.scheme}:// protocol under uri/readUTF8`
+        logError('Agent', 'unsupported operation', errorMessage)
+        throw new Error(errorMessage)
+    }
+
     private maybeExtension: ExtensionObjects | undefined
 
     public async provide(extension: ExtensionObjects): Promise<vscode.Disposable> {

vscode/src/chat/chat-view/ChatController.ts

@@ -70,6 +70,7 @@ import type { RemoteRepoPicker } from '../../context/repo-picker'
 import { resolveContextItems } from '../../editor/utils/editor-context'
 import type { VSCodeEditor } from '../../editor/vscode-editor'
 import type { ExtensionClient } from '../../extension-client'
+import type { ClientCapabilities } from '../../jsonrpc/agent-protocol'
 import type { LocalEmbeddingsController } from '../../local-context/local-embeddings'
 import type { SymfRunner } from '../../local-context/symf'
 import { logDebug } from '../../log'
@@ -1824,32 +1825,82 @@ export async function addWebviewViewHTML(
         return
     }
     const config = extensionClient.capabilities?.webviewNativeConfig
-    const webviewPath = config?.rootDir
-        ? vscode.Uri.parse(config?.rootDir, true)
-        : vscode.Uri.joinPath(extensionUri, 'dist', 'webviews')
-    // Create Webview using vscode/index.html
-    const root = vscode.Uri.joinPath(webviewPath, 'index.html')
-    const bytes = await vscode.workspace.fs.readFile(root)
-    const decoded = new TextDecoder('utf-8').decode(bytes)
-    const resources = view.webview.asWebviewUri(webviewPath)
+    const baseUri = getWebviewBase(extensionUri, extensionClient.capabilities)
+    const html = await loadWebviewIndex(extensionClient, baseUri)
+
+    let resources: vscode.Uri | undefined
+    if (config?.assetLoader !== 'webviewasset') {
+        resources = view.webview.asWebviewUri(baseUri)
+    }
+
+    view.webview.html = transformHTML({ html, resources, config, cspSource: view.webview.cspSource })
+}
 
+interface TransformHTMLOptions {
+    html: string
+    resources?: vscode.Uri
+    config?: ClientCapabilities['webviewNativeConfig']
+    cspSource: string
+}
+
+const transformHTML = ({ html, resources, config, cspSource }: TransformHTMLOptions): string => {
     // This replace variables from the vscode/dist/index.html with webview info
     // 1. Update URIs to load styles and scripts into webview (eg. path that starts with ./)
+    if (resources) {
+        html = html.replaceAll('./', `${resources.toString()}/`)
+    }
     // 2. Update URIs for content security policy to only allow specific scripts to be run
-    view.webview.html = decoded
-        .replaceAll('./', `${resources.toString()}/`)
-        .replaceAll("'self'", view.webview.cspSource)
-        .replaceAll('{cspSource}', view.webview.cspSource)
+    html = html.replaceAll("'self'", cspSource).replaceAll('{cspSource}', cspSource)
 
     // If a script or style is injected, replace the placeholder with the script or style
     // and drop the content-security-policy meta tag which prevents inline scripts and styles
     if (config?.injectScript || config?.injectStyle) {
-        // drop all text betweeb <-- START CSP --> and <-- END CSP -->
-        view.webview.html = decoded
+        // drop all text between <-- START CSP --> and <-- END CSP -->
+        html = html
             .replace(/<-- START CSP -->.*<!-- END CSP -->/s, '')
             .replaceAll('/*injectedScript*/', config?.injectScript ?? '')
             .replaceAll('/*injectedStyle*/', config?.injectStyle ?? '')
     }
+
+    return html
+}
+
+const loadWebviewIndex = async (
+    extensionClient: ExtensionClient,
+    basePath: vscode.Uri
+): Promise<string> => {
+    const index = vscode.Uri.joinPath(basePath, 'index.html')
+    if (
+        extensionClient.capabilities?.uriSchemeLoaders?.includes(index.scheme) &&
+        extensionClient.readUriUTF8
+    ) {
+        const utf8 = await extensionClient.readUriUTF8(index)
+        if (!utf8) {
+            throw new Error('Failed to load webview asset: ' + index.toString())
+        }
+        return utf8
+    }
+
+    const bytes = await vscode.workspace.fs.readFile(index)
+    return new TextDecoder('utf-8').decode(bytes)
+}
+
+const getWebviewBase = (extensionUri: vscode.Uri, capabilities?: ClientCapabilities): vscode.Uri => {
+    const config = capabilities?.webviewNativeConfig
+
+    // When the client has specified a webview asset loader, and they support webviewasset://,
+    // we will delegate to the client to handle the URI.
+    if (
+        config?.assetLoader === 'webviewasset' &&
+        capabilities?.uriSchemeLoaders?.includes('webviewasset')
+    ) {
+        return vscode.Uri.from({ scheme: 'webviewasset', path: '/' })
+    }
+    // Otherwise, we will use load the webview asset from the file system, either
+    // from a specfic directory provided by the client, or from the default location.
+    return config?.rootDir
+        ? vscode.Uri.parse(config?.rootDir, true)
+        : vscode.Uri.joinPath(extensionUri, 'dist', 'webviews')
 }
 
 // This is the manual ordering of the different retrieved and explicit context sources

vscode/src/extension-client.ts

@@ -38,6 +38,8 @@ export interface ExtensionClient {
      */
     openNewDocument(workspace: typeof vscode.workspace, uri: Uri): Thenable<TextDocument | undefined>
 
+    readUriUTF8?(uri: Uri): Thenable<string | undefined>
+
     get clientName(): string
     get clientVersion(): string
     get capabilities(): ClientCapabilities | undefined

vscode/src/jsonrpc/agent-protocol.ts

@@ -343,6 +343,8 @@ export type ServerRequests = {
     ]
     'workspace/edit': [WorkspaceEditParams, boolean]
 
+    'uri/readUTF8': [{ uri: string }, { text: string }]
+
     'secrets/get': [{ key: string }, string | null | undefined]
     'secrets/store': [{ key: string; value: string }, null | undefined]
     'secrets/delete': [{ key: string }, null | undefined]
@@ -648,22 +650,32 @@ export interface ClientCapabilities {
     // which effectively means both sidebar and custom editor chat views are supported.
     // Defaults to 'agentic'.
     webview?: 'agentic' | 'native' | undefined | null
+
+    // Custom URI prefixes that this client supports. If the agent is asked to open a URI
+    // with one of these prefixes, it will delegate to the client to handle the URI.
+    uriSchemeLoaders?: string[] | undefined | null
+
     // If webview === 'native', describes how the client has configured webview resources.
     // cspSource is passed to the extension as the Webview cspSource property.
     // webviewBundleServingPrefix is prepended to resource paths under 'dist' in
     // asWebviewUri (note, multiple prefixes are not yet implemented.)
     // Set the view to 'single' when client only support single chat view, e.g. sidebar chat.
-    webviewNativeConfig?:
-        | {
-              view: 'multiple' | 'single'
-              cspSource: string
-              webviewBundleServingPrefix: string
-              rootDir?: string | undefined | null
-              injectScript?: string | undefined | null
-              injectStyle?: string | undefined | null
-          }
-        | undefined
-        | null
+    webviewNativeConfig?: WebviewNativeConfig | undefined | null
+}
+
+export interface WebviewNativeConfig {
+    view: 'multiple' | 'single'
+    cspSource: string
+    // if assetLoader is 'webviewasset', the client must implement the
+    // webviewasset:// protocol. The agent will call into the extension client
+    // to resolve the webview. If assetLoader is 'fs' or undefined, the
+    // agent will attempt to load the asset from the file system at
+    // (rootDir ?? codyPaths())/dist/webviews.
+    assetLoader?: 'fs' | 'webviewasset' | undefined | null
+    webviewBundleServingPrefix: string
+    rootDir?: string | undefined | null
+    injectScript?: string | undefined | null
+    injectStyle?: string | undefined | null
 }
 
 export interface ServerInfo {