Agent: add new secrets capability to implement secret storage (#5348)

Previously, the agent only supported stateless secret storage, where the
agent server stored secrets in a temporary hashmap that was lost
whenever the agent started. Now, clients can optionally declare that
they're able to store secrets using the `secrets: 'client-managed'`
capability. With this new capability, client can
store/retrieve/delete/change secrets using the new JSON-RPC methods:

* `secrets/get`
* `secrets/store`
* `secrets/delete`
* `secrets/didChange`

Here's the PR moving to client-managed secrets for the Eclipse plugin,
which allowed us to delete 600 lines of native UI code 😮
sourcegraph/eclipse#54


## Test plan

This PR changes `TestClient` to use client-managed secrets, so we're
stressing this new code path in all the integration tests by default.

<!-- Required. See
https://docs-legacy.sourcegraph.com/dev/background-information/testing_principles.
-->

## Changelog

<!-- OPTIONAL; info at
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c
-->

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

@@ -18,6 +18,7 @@ data class ClientCapabilities(
   val codeActions: CodeActionsEnum? = null, // Oneof: none, enabled
   val webviewMessages: WebviewMessagesEnum? = null, // Oneof: object-encoded, string-encoded
   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,
 ) {
@@ -92,6 +93,11 @@ data class ClientCapabilities(
     @SerializedName("client-managed") `Client-managed`,
   }
 
+  enum class SecretsEnum {
+    @SerializedName("stateless") Stateless,
+    @SerializedName("client-managed") `Client-managed`,
+  }
+
   enum class WebviewEnum {
     @SerializedName("agentic") Agentic,
     @SerializedName("native") Native,

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

@@ -22,6 +22,12 @@ interface CodyAgentClient {
   fun textDocument_show(params: TextDocument_ShowParams): CompletableFuture<Boolean>
   @JsonRequest("workspace/edit")
   fun workspace_edit(params: WorkspaceEditParams): CompletableFuture<Boolean>
+  @JsonRequest("secrets/get")
+  fun secrets_get(params: Secrets_GetParams): CompletableFuture<String?>
+  @JsonRequest("secrets/store")
+  fun secrets_store(params: Secrets_StoreParams): CompletableFuture<String?>
+  @JsonRequest("secrets/delete")
+  fun secrets_delete(params: Secrets_DeleteParams): CompletableFuture<String?>
   @JsonRequest("env/openExternal")
   fun env_openExternal(params: Env_OpenExternalParams): CompletableFuture<Boolean>
 

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

@@ -178,4 +178,6 @@ interface CodyAgentServer {
   fun progress_cancel(params: Progress_CancelParams)
   @JsonNotification("webview/didDisposeNative")
   fun webview_didDisposeNative(params: Webview_DidDisposeNativeParams)
+  @JsonNotification("secrets/didChange")
+  fun secrets_didChange(params: Secrets_DidChangeParams)
 }

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

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

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

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

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

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

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

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

agent/src/AgentSecretStorage.ts

@@ -0,0 +1,36 @@
+import type * as vscode from 'vscode'
+import { emptyEvent } from '../../vscode/src/testutils/emptyEvent'
+import type { MessageHandler } from './jsonrpc-alias'
+
+export class AgentStatelessSecretStorage implements vscode.SecretStorage {
+    private readonly inMemorySecretStorageMap = new Map<string, string>()
+    public get(key: string): Thenable<string | undefined> {
+        return Promise.resolve(this.inMemorySecretStorageMap.get(key))
+    }
+    public store(key: string, value: string): Thenable<void> {
+        this.inMemorySecretStorageMap.set(key, value)
+        return Promise.resolve()
+    }
+    public delete(key: string): Thenable<void> {
+        this.inMemorySecretStorageMap.delete(key)
+        return Promise.resolve()
+    }
+    onDidChange: vscode.Event<vscode.SecretStorageChangeEvent> = emptyEvent()
+}
+
+export class AgentClientManagedSecretStorage implements vscode.SecretStorage {
+    constructor(
+        private readonly agent: MessageHandler,
+        public readonly onDidChange: vscode.Event<vscode.SecretStorageChangeEvent>
+    ) {}
+    public async get(key: string): Promise<string | undefined> {
+        const result = await this.agent.request('secrets/get', { key })
+        return result ?? undefined
+    }
+    public async store(key: string, value: string): Promise<void> {
+        await this.agent.request('secrets/store', { key, value })
+    }
+    public async delete(key: string): Promise<void> {
+        await this.agent.request('secrets/delete', { key })
+    }
+}

agent/src/TestClient.ts

@@ -25,6 +25,7 @@ import {
     TESTING_CREDENTIALS,
     type TestingCredentials,
 } from '../../vscode/src/testutils/testing-credentials'
+import { AgentStatelessSecretStorage } from './AgentSecretStorage'
 import { AgentTextDocument } from './AgentTextDocument'
 import { AgentWorkspaceDocuments } from './AgentWorkspaceDocuments'
 import { allClientCapabilitiesEnabled } from './allClientCapabilitiesEnabled'
@@ -132,6 +133,7 @@ export function buildAgentBinary(): void {
 }
 
 export class TestClient extends MessageHandler {
+    private secrets = new AgentStatelessSecretStorage()
     private extensionConfigurationDuringInitialization: ExtensionConfiguration | undefined
     public static create({ bin = 'node', ...params }: TestClientParams): TestClient {
         buildAgentBinary()
@@ -240,6 +242,17 @@ export class TestClient extends MessageHandler {
                 message: {},
             })
         })
+        this.registerRequest('secrets/get', async ({ key }) => {
+            return this.secrets.get(key) ?? null
+        })
+        this.registerRequest('secrets/store', async ({ key, value }) => {
+            await this.secrets.store(key, value)
+            return null
+        })
+        this.registerRequest('secrets/delete', async ({ key }) => {
+            await this.secrets.delete(key)
+            return null
+        })
         this.registerRequest('window/showMessage', params => {
             if (this.params.onWindowRequest) {
                 return this.params.onWindowRequest(params)

agent/src/agent.ts

@@ -60,9 +60,9 @@ import { IndentationBasedFoldingRangeProvider } from '../../vscode/src/lsp/foldi
 import type { FixupActor, FixupFileCollection } from '../../vscode/src/non-stop/roles'
 import type { FixupControlApplicator } from '../../vscode/src/non-stop/strategies'
 import { AgentWorkspaceEdit } from '../../vscode/src/testutils/AgentWorkspaceEdit'
-import { emptyEvent } from '../../vscode/src/testutils/emptyEvent'
 import { AgentFixupControls } from './AgentFixupControls'
 import { AgentProviders } from './AgentProviders'
+import { AgentClientManagedSecretStorage, AgentStatelessSecretStorage } from './AgentSecretStorage'
 import { AgentWebviewPanel, AgentWebviewPanels } from './AgentWebviewPanel'
 import { AgentWorkspaceDocuments } from './AgentWorkspaceDocuments'
 import { registerNativeWebviewHandlers, resolveWebviewView } from './NativeWebview'
@@ -93,8 +93,6 @@ import type {
 import * as vscode_shim from './vscode-shim'
 import { vscodeLocation, vscodeRange } from './vscode-type-converters'
 
-const inMemorySecretStorageMap = new Map<string, string>()
-
 /** The VS Code extension's `activate` function. */
 type ExtensionActivate = (
     context: vscode.ExtensionContext,
@@ -140,7 +138,8 @@ export async function initializeVscodeExtension(
     workspaceRoot: vscode.Uri,
     extensionActivate: ExtensionActivate,
     extensionClient: ExtensionClient,
-    globalState: AgentGlobalState
+    globalState: AgentGlobalState,
+    secrets: vscode.SecretStorage
 ): Promise<void> {
     const paths = codyPaths()
     const extensionPath = paths.config
@@ -161,19 +160,7 @@ export async function initializeVscodeExtension(
         globalState,
         logUri: vscode.Uri.file(paths.log),
         logPath: paths.log,
-        secrets: {
-            onDidChange: emptyEvent(),
-            get(key) {
-                return Promise.resolve(inMemorySecretStorageMap.get(key))
-            },
-            store(key, value) {
-                inMemorySecretStorageMap.set(key, value)
-                return Promise.resolve()
-            },
-            delete() {
-                return Promise.resolve()
-            },
-        },
+        secrets,
         storageUri: vscode.Uri.file(paths.data),
         subscriptions: [],
 
@@ -347,6 +334,7 @@ export class Agent extends MessageHandler implements ExtensionClient {
             })
         },
     })
+    private secretsDidChange = new vscode.EventEmitter<vscode.SecretStorageChangeEvent>()
 
     public webPanels = new AgentWebviewPanels()
     public webviewViewProviders = new Map<string, vscode.WebviewViewProvider>()
@@ -432,11 +420,17 @@ export class Agent extends MessageHandler implements ExtensionClient {
                   })
 
             try {
+                const secrets =
+                    clientInfo.capabilities?.secrets === 'client-managed'
+                        ? new AgentClientManagedSecretStorage(this, this.secretsDidChange.event)
+                        : new AgentStatelessSecretStorage()
+
                 await initializeVscodeExtension(
                     this.workspace.workspaceRootUri,
                     params.extensionActivate,
                     this,
-                    this.globalState
+                    this.globalState,
+                    secrets
                 )
 
                 this.authenticationPromise = clientInfo.extensionConfiguration
@@ -1370,6 +1364,9 @@ export class Agent extends MessageHandler implements ExtensionClient {
                 return null
             }
         )
+        this.registerNotification('secrets/didChange', async ({ key }) => {
+            this.secretsDidChange.fire({ key })
+        })
 
         this.registerAuthenticatedRequest('featureFlags/getFeatureFlag', async ({ flagName }) => {
             return featureFlagProvider.evaluateFeatureFlag(

agent/src/allClientCapabilitiesEnabled.ts

@@ -10,4 +10,5 @@ export const allClientCapabilitiesEnabled: ClientCapabilities = {
     showWindowMessage: 'request',
     ignore: 'enabled',
     codeActions: 'enabled',
+    secrets: 'client-managed',
 }

agent/src/bfg/BfgRetriever.test.ts

@@ -14,6 +14,7 @@ import { getCurrentDocContext } from '../../../vscode/src/completions/get-curren
 import { initTreeSitterParser } from '../../../vscode/src/completions/test-helpers'
 import { defaultVSCodeExtensionClient } from '../../../vscode/src/extension-client'
 import { activate } from '../../../vscode/src/extension.node'
+import { AgentStatelessSecretStorage } from '../AgentSecretStorage'
 import { initializeVscodeExtension, newEmbeddedAgentClient } from '../agent'
 import { AgentGlobalState } from '../global-state/AgentGlobalState'
 import * as vscode_shim from '../vscode-shim'
@@ -41,7 +42,8 @@ describe('BfgRetriever', async () => {
             vscode.Uri.file(process.cwd()),
             activate,
             defaultVSCodeExtensionClient(),
-            new AgentGlobalState('vscode')
+            new AgentGlobalState('vscode'),
+            new AgentStatelessSecretStorage()
         )
 
         if (shouldCreateGitDir) {

vscode/src/jsonrpc/agent-protocol.ts

@@ -343,6 +343,10 @@ export type ServerRequests = {
     ]
     'workspace/edit': [WorkspaceEditParams, boolean]
 
+    'secrets/get': [{ key: string }, string | null | undefined]
+    'secrets/store': [{ key: string; value: string }, null | undefined]
+    'secrets/delete': [{ key: string }, null | undefined]
+
     // TODO: Add VSCode support for registerWebviewPanelSerializer.
 
     'env/openExternal': [{ uri: string }, boolean]
@@ -415,6 +419,8 @@ export type ClientNotifications = {
     // Consequently they have their own dispose notification. c.f.
     // webview/dispose client request.
     'webview/didDisposeNative': [{ handle: string }]
+
+    'secrets/didChange': [{ key: string }]
 }
 
 // ================
@@ -623,6 +629,12 @@ export interface ClientCapabilities {
     // JSON-RPC request to handle the saving of the client state. This is needed to safely share state
     // between concurrent agent processes (assuming there is one IDE client process managing multiple agent processes).
     globalState?: 'stateless' | 'server-managed' | 'client-managed' | undefined | null
+
+    // Secrets controls how the agent should handle storing secrets.
+    // - Stateless: the secrets are not persisted between agent processes.
+    // - Client managed: the client must implement the 'secrets/get',
+    // 'secrets/store', and 'secrets/delete' requests.
+    secrets?: 'stateless' | 'client-managed' | undefined | null
     // Whether the client supports the VSCode WebView API. If 'agentic', uses
     // AgentWebViewPanel which just delegates bidirectional postMessage over
     // the Agent protocol. If 'native', implements a larger subset of the VSCode