Autocomplete: Enable keepalive agent for Node (#868)

We found another small way to improve overall latencies: Making sure we
keep the TCP connection alive and reducing the need for subsequent
SSL/TLS handshakes.

In Node, reusing TCP connections is not the default behavior (like it is
for example on our Go backends). To fix this, we have to make sure we
create a custom network agent that is configured to keep connections
allive.

Since this is Node only, we are using a mutable ref objects to not run
any of those code paths on the web branches.

c.f.
https://sourcegraph.slack.com/archives/C05497E9MDW/p1693387804620389

## Test plan

Tested using the local completion CLI script

We found about a 5% latency improvement:


![image](https://github.com/sourcegraph/cody/assets/458591/d667d71a-3c04-444d-b6c8-08ea3cca70f4)


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

---------

Co-authored-by: Valery Bugakov <skymk1@gmail.com>

vscode/CHANGELOG.md

@@ -23,6 +23,7 @@ Starting from `0.2.0`, Cody is using `major.EVEN_NUMBER.patch` for release versi
 
 ### Changed
 
+- Autocomplete on VS Code desktop instances now reuses TCP connections to reduce latency. [pull/868](https://github.com/sourcegraph/cody/pull/868)
 - Errors are now always logged to the output console, even if the debug mode is not enabled. [pull/851](https://github.com/sourcegraph/cody/pull/851)
 
 ## [0.8.0]

vscode/src/completions/client.ts

@@ -9,6 +9,8 @@ import type {
 } from '@sourcegraph/cody-shared/src/sourcegraph-api/completions/types'
 import { NetworkError, RateLimitError } from '@sourcegraph/cody-shared/src/sourcegraph-api/errors'
 
+import { fetch } from '../fetch'
+
 export type CodeCompletionsParams = Omit<CompletionParameters, 'fast'>
 
 export interface CodeCompletionsClient {
@@ -44,6 +46,9 @@ export function createClient(
                 : [false, false]
 
             const headers = new Headers(config.customHeaders)
+            // Force HTTP connection reuse to reduce latency.
+            // c.f. https://github.com/microsoft/vscode/issues/173861
+            headers.set('Connection', 'keep-alive')
             if (config.accessToken) {
                 headers.set('Authorization', `token ${config.accessToken}`)
             }

vscode/src/extension.node.ts

@@ -11,6 +11,7 @@ import { CommandsController } from './custom-prompts/CommandsController'
 import { ExtensionApi } from './extension-api'
 import { activate as activateCommon } from './extension.common'
 import { VSCODE_WEB_RECIPES } from './extension.web'
+import { initializeNetworkAgent } from './fetch.node'
 import { FilenameContextFetcher } from './local-context/filename-context-fetcher'
 import { LocalKeywordContextFetcher } from './local-context/local-keyword-context-fetcher'
 import { SymfRunner } from './local-context/symf'
@@ -21,6 +22,8 @@ import { getRgPath } from './rg'
  * (Node.js/Electron).
  */
 export function activate(context: vscode.ExtensionContext): ExtensionApi {
+    initializeNetworkAgent()
+
     return activateCommon(context, {
         getRgPath,
         createCommandsController: (...args) => new CommandsController(...args),

vscode/src/fetch.node.ts

@@ -0,0 +1,66 @@
+import http from 'http'
+import https from 'https'
+
+import { agent } from './fetch'
+
+// The path to the exported class can be found in the npm contents
+// https://www.npmjs.com/package/@vscode/proxy-agent?activeTab=code
+const nodeModules = '_VSCODE_NODE_MODULES'
+const proxyAgentPath = '@vscode/proxy-agent/out/agent'
+const proxyAgent = 'PacProxyAgent'
+
+export function initializeNetworkAgent(): void {
+    /**
+     * We use keepAlive agents here to avoid excessive SSL/TLS handshakes for autocomplete requests.
+     */
+    const httpAgent = new http.Agent({ keepAlive: true, keepAliveMsecs: 60000 })
+    const httpsAgent = new https.Agent({ keepAlive: true, keepAliveMsecs: 60000 })
+
+    const customAgent = ({ protocol }: Pick<URL, 'protocol'>): http.Agent => {
+        if (protocol === 'http:') {
+            return httpAgent
+        }
+        return httpsAgent
+    }
+
+    agent.current = customAgent
+
+    /**
+     * This works around an issue in the default VS Code proxy agent code. When `http.proxySupport`
+     * is set to its default value and no proxy setting is being used, the proxy library does not
+     * properly reuse the agent set on the http(s) method and is instead always using a new agent
+     * per request.
+     *
+     * To work around this, we patch the default proxy agent method and overwrite the
+     * `originalAgent` value before invoking it for requests that want to keep their connection
+     * alive (as indicated by the `Connection: keep-alive` header).
+     *
+     * c.f. https://github.com/microsoft/vscode/issues/173861
+     */
+    try {
+        const PacProxyAgent = (globalThis as any)?.[nodeModules]?.[proxyAgentPath]?.[proxyAgent] ?? undefined
+        if (PacProxyAgent) {
+            const originalConnect = PacProxyAgent.prototype.connect
+            // Patches the implementation defined here:
+            // https://github.com/microsoft/vscode-proxy-agent/blob/d340b9d34684da494d6ebde3bcd18490a8bbd071/src/agent.ts#L53
+            PacProxyAgent.prototype.connect = function (req: http.ClientRequest, opts: { protocol: string }): any {
+                try {
+                    const connectionHeader = req.getHeader('connection')
+                    if (
+                        connectionHeader === 'keep-alive' ||
+                        (Array.isArray(connectionHeader) && connectionHeader.includes('keep-alive'))
+                    ) {
+                        this.opts.originalAgent = customAgent(opts)
+                        return originalConnect.call(this, req, opts)
+                    }
+                    return originalConnect.call(this, req, opts)
+                } catch {
+                    return originalConnect.call(this, req, opts)
+                }
+            }
+        }
+    } catch (error) {
+        // Ignore any errors in the patching logic
+        void error
+    }
+}

vscode/src/fetch.ts

@@ -0,0 +1,26 @@
+import type { Agent } from 'http'
+
+/**
+ * By hard-requiring isomorphic-fetch, we ensure that even in newer Node environments that include
+ * `fetch` by default, we still use the `node-fetch` polyfill and have access to the networking code
+ */
+import isomorphicFetch from 'isomorphic-fetch'
+
+/**
+ * In node environments, it might be necessary to set up a custom agent to control the network
+ * requests being made.
+ *
+ * To do this, we have a mutable agent variable that can be set to an instance of `http.Agent` or
+ * `https.Agent` (depending on the protocol of the URL) but that will be kept undefined for web
+ * environments.
+ *
+ * Agent is a mutable ref so that we can override it from `fetch.node.ts`
+ */
+export const agent: { current: ((url: URL) => Agent) | undefined } = { current: undefined }
+
+export function fetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
+    return isomorphicFetch(input, {
+        ...init,
+        agent: agent.current,
+    } as RequestInit)
+}

vscode/test/completions/run-code-completions-on-dataset.ts

@@ -20,6 +20,7 @@ import { ProviderConfig } from '../../src/completions/providers/provider'
 import { InlineCompletionItemProvider } from '../../src/completions/vscodeInlineCompletionItemProvider'
 import { getFullConfig } from '../../src/configuration'
 import { configureExternalServices } from '../../src/external-services'
+import { initializeNetworkAgent } from '../../src/fetch.node'
 import { InMemorySecretStorage } from '../../src/services/SecretStorageProvider'
 import { wrapVSCodeTextDocument } from '../../src/testutils/textDocument'
 
@@ -38,6 +39,8 @@ const dummyFeatureFlagProvider = new FeatureFlagProvider(
     })
 )
 
+initializeNetworkAgent()
+
 async function initCompletionsProvider(context: GetContextResult): Promise<InlineCompletionItemProvider> {
     const secretStorage = new InMemorySecretStorage()
     await secretStorage.store('cody.access-token', ENVIRONMENT_CONFIG.SOURCEGRAPH_ACCESS_TOKEN)