From d5cfa4616bbff23734a8b0950f5eef03c713bdfc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=93lafur=20P=C3=A1ll=20Geirsson?= Date: Tue, 27 Aug 2024 20:56:42 +0200 Subject: [PATCH] Agent: add new `secrets` capability to implement secret storage (#5348) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 😮 https://github.com/sourcegraph/eclipse/pull/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. ## Changelog --- .../protocol_generated/ClientCapabilities.kt | 6 ++++ .../protocol_generated/CodyAgentClient.kt | 6 ++++ .../protocol_generated/CodyAgentServer.kt | 2 ++ .../Secrets_DeleteParams.kt | 7 ++++ .../Secrets_DidChangeParams.kt | 7 ++++ .../protocol_generated/Secrets_GetParams.kt | 7 ++++ .../protocol_generated/Secrets_StoreParams.kt | 8 +++++ agent/src/AgentSecretStorage.ts | 36 +++++++++++++++++++ agent/src/TestClient.ts | 13 +++++++ agent/src/agent.ts | 33 ++++++++--------- agent/src/allClientCapabilitiesEnabled.ts | 1 + agent/src/bfg/BfgRetriever.test.ts | 4 ++- vscode/src/jsonrpc/agent-protocol.ts | 12 +++++++ 13 files changed, 123 insertions(+), 19 deletions(-) create mode 100644 agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DeleteParams.kt create mode 100644 agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DidChangeParams.kt create mode 100644 agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_GetParams.kt create mode 100644 agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_StoreParams.kt create mode 100644 agent/src/AgentSecretStorage.ts diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ClientCapabilities.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ClientCapabilities.kt index 0961dda77d8c..115fdcd4b0a2 100644 --- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ClientCapabilities.kt +++ b/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, diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt index fb4a12950781..4caf7e7be9df 100644 --- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt +++ b/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 @JsonRequest("workspace/edit") fun workspace_edit(params: WorkspaceEditParams): CompletableFuture + @JsonRequest("secrets/get") + fun secrets_get(params: Secrets_GetParams): CompletableFuture + @JsonRequest("secrets/store") + fun secrets_store(params: Secrets_StoreParams): CompletableFuture + @JsonRequest("secrets/delete") + fun secrets_delete(params: Secrets_DeleteParams): CompletableFuture @JsonRequest("env/openExternal") fun env_openExternal(params: Env_OpenExternalParams): CompletableFuture diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentServer.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentServer.kt index c475284b5a47..a4ea0fcb37e2 100644 --- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentServer.kt +++ b/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) } diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DeleteParams.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DeleteParams.kt new file mode 100644 index 000000000000..29bf5fdb7ad5 --- /dev/null +++ b/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, +) + diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DidChangeParams.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_DidChangeParams.kt new file mode 100644 index 000000000000..b69adf422c53 --- /dev/null +++ b/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, +) + diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_GetParams.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_GetParams.kt new file mode 100644 index 000000000000..84cb542d2f40 --- /dev/null +++ b/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, +) + diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_StoreParams.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Secrets_StoreParams.kt new file mode 100644 index 000000000000..d2f7ec0f8433 --- /dev/null +++ b/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, +) + diff --git a/agent/src/AgentSecretStorage.ts b/agent/src/AgentSecretStorage.ts new file mode 100644 index 000000000000..fa5039cf09bd --- /dev/null +++ b/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() + public get(key: string): Thenable { + return Promise.resolve(this.inMemorySecretStorageMap.get(key)) + } + public store(key: string, value: string): Thenable { + this.inMemorySecretStorageMap.set(key, value) + return Promise.resolve() + } + public delete(key: string): Thenable { + this.inMemorySecretStorageMap.delete(key) + return Promise.resolve() + } + onDidChange: vscode.Event = emptyEvent() +} + +export class AgentClientManagedSecretStorage implements vscode.SecretStorage { + constructor( + private readonly agent: MessageHandler, + public readonly onDidChange: vscode.Event + ) {} + public async get(key: string): Promise { + const result = await this.agent.request('secrets/get', { key }) + return result ?? undefined + } + public async store(key: string, value: string): Promise { + await this.agent.request('secrets/store', { key, value }) + } + public async delete(key: string): Promise { + await this.agent.request('secrets/delete', { key }) + } +} diff --git a/agent/src/TestClient.ts b/agent/src/TestClient.ts index e8f47023d424..eca2eb07eaaa 100644 --- a/agent/src/TestClient.ts +++ b/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) diff --git a/agent/src/agent.ts b/agent/src/agent.ts index fb9cfe630564..f1501756c017 100644 --- a/agent/src/agent.ts +++ b/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() - /** 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 { 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() public webPanels = new AgentWebviewPanels() public webviewViewProviders = new Map() @@ -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( diff --git a/agent/src/allClientCapabilitiesEnabled.ts b/agent/src/allClientCapabilitiesEnabled.ts index 8d1800b37893..dbd9d6e4350f 100644 --- a/agent/src/allClientCapabilitiesEnabled.ts +++ b/agent/src/allClientCapabilitiesEnabled.ts @@ -10,4 +10,5 @@ export const allClientCapabilitiesEnabled: ClientCapabilities = { showWindowMessage: 'request', ignore: 'enabled', codeActions: 'enabled', + secrets: 'client-managed', } diff --git a/agent/src/bfg/BfgRetriever.test.ts b/agent/src/bfg/BfgRetriever.test.ts index bd772135776a..135ff923d2d6 100644 --- a/agent/src/bfg/BfgRetriever.test.ts +++ b/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) { diff --git a/vscode/src/jsonrpc/agent-protocol.ts b/vscode/src/jsonrpc/agent-protocol.ts index 8e19defc7c4c..6425c7ff937c 100644 --- a/vscode/src/jsonrpc/agent-protocol.ts +++ b/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