diff --git a/src/services/code-index/embedders/__tests__/bedrock.spec.ts b/src/services/code-index/embedders/__tests__/bedrock.spec.ts index dfa9544715..f089fece39 100644 --- a/src/services/code-index/embedders/__tests__/bedrock.spec.ts +++ b/src/services/code-index/embedders/__tests__/bedrock.spec.ts @@ -3,6 +3,7 @@ import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedroc import { BedrockEmbedder } from "../bedrock" import { MAX_ITEM_TOKENS, INITIAL_RETRY_DELAY_MS } from "../../constants" +import { createProxyRoutingRequestHandler } from "../../../../utils/networkProxy" import { clearAllMocks } from "../../../../test-utils/reset" @@ -26,6 +27,10 @@ vitest.mock("@aws-sdk/credential-providers", () => ({ fromIni: vitest.fn().mockReturnValue(Promise.resolve({})), })) +vitest.mock("../../../../utils/networkProxy", () => ({ + createProxyRoutingRequestHandler: vitest.fn().mockReturnValue(undefined), +})) + // Mock TelemetryService vitest.mock("@roo-code/telemetry", () => ({ TelemetryService: { @@ -92,6 +97,11 @@ describe("BedrockEmbedder", () => { }) describe("constructor", () => { + afterEach(() => { + // clearAllMocks() keeps implementations, so a stub would leak into later tests. + vitest.mocked(createProxyRoutingRequestHandler).mockReturnValue(undefined) + }) + it("should initialize with provided region, profile and model", () => { expect(embedder.embedderInfo.name).toBe("bedrock") }) @@ -114,6 +124,34 @@ describe("BedrockEmbedder", () => { }), ) }) + + it("should route requests through the proxy-aware handler when one is built", () => { + const handler = { + handle: vitest.fn(), + updateHttpClientConfig: vitest.fn(), + httpHandlerConfigs: vitest.fn(), + destroy: vitest.fn(), + } + vitest.mocked(createProxyRoutingRequestHandler).mockReturnValue(handler) + + new BedrockEmbedder("us-east-1", "test-profile", "amazon.titan-embed-text-v2:0") + + const clientConfig = vitest.mocked(BedrockRuntimeClient).mock.calls.at(-1)?.[0] + expect(clientConfig?.requestHandler).toBe(handler) + // Pinning an endpoint here would break FIPS, dualstack and non-default partitions, + // and would take the routing decision away from the handler. + expect(clientConfig).not.toHaveProperty("endpoint") + }) + + it("should keep the client default handler when no proxy is configured", () => { + vitest.mocked(createProxyRoutingRequestHandler).mockReturnValue(undefined) + + new BedrockEmbedder("us-east-1", "test-profile", "amazon.titan-embed-text-v2:0") + + expect(BedrockRuntimeClient).toHaveBeenLastCalledWith( + expect.not.objectContaining({ requestHandler: expect.anything() }), + ) + }) }) describe("createEmbeddings", () => { diff --git a/src/services/code-index/embedders/bedrock.ts b/src/services/code-index/embedders/bedrock.ts index 833d4bf8f8..2473346562 100644 --- a/src/services/code-index/embedders/bedrock.ts +++ b/src/services/code-index/embedders/bedrock.ts @@ -1,5 +1,6 @@ import { BedrockRuntimeClient, InvokeModelCommand, InvokeModelCommandInput } from "@aws-sdk/client-bedrock-runtime" import { fromIni, fromNodeProviderChain } from "@aws-sdk/credential-providers" +import { createProxyRoutingRequestHandler } from "../../../utils/networkProxy" import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces" import { MAX_BATCH_TOKENS, @@ -40,10 +41,16 @@ export class BedrockEmbedder implements IEmbedder { // If profile is specified, use it; otherwise use default credential chain const credentials = this.profile ? fromIni({ profile: this.profile }) : fromNodeProviderChain() + // Behind a corporate proxy, Node resolves DNS locally and Bedrock endpoints fail with + // ENOTFOUND. The handler tunnels through the proxy with CONNECT so the proxy resolves + // the hostname, and connects directly to the destinations NO_PROXY excludes. + const requestHandler = createProxyRoutingRequestHandler() + this.bedrockClient = new BedrockRuntimeClient({ userAgentAppId: `ZooCode#${Package.version}`, region: this.region, credentials, + ...(requestHandler && { requestHandler }), }) this.defaultModelId = modelId || getDefaultModelId("bedrock") diff --git a/src/utils/__tests__/networkProxy.spec.ts b/src/utils/__tests__/networkProxy.spec.ts index 94e91cd990..41c85a10e9 100644 --- a/src/utils/__tests__/networkProxy.spec.ts +++ b/src/utils/__tests__/networkProxy.spec.ts @@ -1,11 +1,35 @@ import * as vscode from "vscode" -import { initializeNetworkProxy, getProxyConfig, isProxyEnabled, isDebugMode, getSystemProxyUrl } from "../networkProxy" +import { NodeHttpHandler } from "@smithy/node-http-handler" +import { HttpProxyAgent } from "http-proxy-agent" +import { HttpsProxyAgent } from "https-proxy-agent" +import { + initializeNetworkProxy, + getProxyConfig, + isProxyEnabled, + isDebugMode, + getSystemProxyUrl, + createProxyRoutingRequestHandler, +} from "../networkProxy" // Mock global-agent vi.mock("global-agent", () => ({ bootstrap: vi.fn(), })) +vi.mock("@smithy/node-http-handler", () => ({ + NodeHttpHandler: vi.fn().mockImplementation(function (options?: unknown) { + return { + options, + handle: vi.fn(), + updateHttpClientConfig: vi.fn(), + httpHandlerConfigs: vi.fn().mockReturnValue({}), + destroy: vi.fn(), + } + }), +})) +vi.mock("http-proxy-agent", () => ({ HttpProxyAgent: vi.fn() })) +vi.mock("https-proxy-agent", () => ({ HttpsProxyAgent: vi.fn() })) + // Mock vscode vi.mock("vscode", () => ({ workspace: { @@ -469,4 +493,118 @@ describe("networkProxy", () => { }) }) }) + + describe("createProxyRoutingRequestHandler", () => { + type Handler = NonNullable> + + // The handler reads only these fields off the request. Building a real HttpRequest would + // pull in @smithy/protocol-http, which is not a direct dependency, so a literal stands in. + const requestTo = (hostname: string, port?: number) => + ({ protocol: "https:", hostname, port }) as unknown as Parameters[0] + + // The two inner handlers are told apart by how they were built: only the proxied one + // receives agents. + const innerHandlers = () => { + const built = vi.mocked(NodeHttpHandler).mock.results.map((r) => r.value) + return { + direct: built.find((h) => h.options === undefined), + proxied: built.find((h) => h.options?.httpsAgent), + } + } + + beforeEach(() => { + vi.clearAllMocks() + delete process.env.HTTPS_PROXY + delete process.env.https_proxy + delete process.env.HTTP_PROXY + delete process.env.http_proxy + delete process.env.NO_PROXY + delete process.env.no_proxy + mockConfig.get.mockReturnValue(undefined) + }) + + it("should return undefined when no proxy is configured", () => { + expect(createProxyRoutingRequestHandler()).toBeUndefined() + expect(NodeHttpHandler).not.toHaveBeenCalled() + }) + + it("should build both proxy agents without keeping connections alive", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + + expect(createProxyRoutingRequestHandler()).toBeDefined() + + // No agent options: an idle connection would outlive the request that opened it. + expect(HttpProxyAgent).toHaveBeenCalledWith("http://proxy.corp:3128") + expect(HttpsProxyAgent).toHaveBeenCalledWith("http://proxy.corp:3128") + const { proxied } = innerHandlers() + expect(proxied?.options?.httpAgent).toBe(vi.mocked(HttpProxyAgent).mock.instances[0]) + expect(proxied?.options?.httpsAgent).toBe(vi.mocked(HttpsProxyAgent).mock.instances[0]) + }) + + it("should send a request through the proxy when NO_PROXY does not cover it", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + process.env.NO_PROXY = "example.com" + + const handler = createProxyRoutingRequestHandler() + const request = requestTo("bedrock-runtime.us-east-1.amazonaws.com") + handler?.handle(request) + + const { direct, proxied } = innerHandlers() + expect(proxied?.handle).toHaveBeenCalledWith(request) + expect(direct?.handle).not.toHaveBeenCalled() + }) + + it("should send a request directly when NO_PROXY covers its host", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + process.env.NO_PROXY = "amazonaws.com" + + const handler = createProxyRoutingRequestHandler() + // The host is only known per request, which is the point of deciding here: indexed + // file contents must not reach a proxy the user excluded. + const request = requestTo("bedrock-runtime.eu-west-1.amazonaws.com") + handler?.handle(request) + + const { direct, proxied } = innerHandlers() + expect(direct?.handle).toHaveBeenCalledWith(request) + expect(proxied?.handle).not.toHaveBeenCalled() + }) + + it("should send every request directly when NO_PROXY is '*'", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + process.env.NO_PROXY = "*" + + const handler = createProxyRoutingRequestHandler() + handler?.handle(requestTo("bedrock-runtime.us-east-1.amazonaws.com")) + + const { direct, proxied } = innerHandlers() + expect(direct?.handle).toHaveBeenCalledOnce() + expect(proxied?.handle).not.toHaveBeenCalled() + }) + + it("should match NO_PROXY against the host when the request carries a port", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + process.env.NO_PROXY = "amazonaws.com" + + const handler = createProxyRoutingRequestHandler() + handler?.handle(requestTo("bedrock-runtime.us-east-1.amazonaws.com", 8443)) + + const { direct, proxied } = innerHandlers() + expect(direct?.handle).toHaveBeenCalledOnce() + expect(proxied?.handle).not.toHaveBeenCalled() + }) + + it("should forward configuration updates and teardown to both routes", () => { + process.env.HTTPS_PROXY = "http://proxy.corp:3128" + + const handler = createProxyRoutingRequestHandler() + handler?.updateHttpClientConfig("requestTimeout", 1234) + handler?.destroy() + + const { direct, proxied } = innerHandlers() + expect(direct?.updateHttpClientConfig).toHaveBeenCalledWith("requestTimeout", 1234) + expect(proxied?.updateHttpClientConfig).toHaveBeenCalledWith("requestTimeout", 1234) + expect(direct?.destroy).toHaveBeenCalledOnce() + expect(proxied?.destroy).toHaveBeenCalledOnce() + }) + }) }) diff --git a/src/utils/networkProxy.ts b/src/utils/networkProxy.ts index 835887334e..e5e75db1f6 100644 --- a/src/utils/networkProxy.ts +++ b/src/utils/networkProxy.ts @@ -11,6 +11,9 @@ */ import * as vscode from "vscode" +import { NodeHttpHandler } from "@smithy/node-http-handler" +import { HttpProxyAgent } from "http-proxy-agent" +import { HttpsProxyAgent } from "https-proxy-agent" import { Package } from "../shared/package" /** @@ -419,6 +422,79 @@ export function getSystemProxyUrl(targetUrl?: string): string | undefined { return undefined } +// Derived from the handler's own signatures so this file needs no @smithy/types dependency. +type HandleParameters = Parameters +type UpdateClientConfigParameters = Parameters + +/** + * The subset of the AWS SDK request-handler contract implemented below. Exported so callers and + * their tests can name the type without reaching for @smithy/types. + */ +export interface ProxyRoutingHandler { + handle(...args: HandleParameters): ReturnType + updateHttpClientConfig(...args: UpdateClientConfigParameters): void + httpHandlerConfigs(): ReturnType + destroy(): void +} + +/** + * AWS SDK request handler that chooses, per request, between the system proxy and a direct + * connection. + * + * The SDK resolves its own endpoint — region, partition, FIPS/dualstack flags, and any endpoint + * override — so the destination is only known once a request has been built. Choosing here is + * what makes NO_PROXY apply to the host actually called, rather than one guessed up front. + * + * The agents take no options, so no connection outlives its request. Both routes are HTTP/1.1: + * a client defaulting to NodeHttp2Handler drops to 1.1 once a proxy is configured, which is what + * the chat provider already does. + */ +class ProxyRoutingRequestHandler implements ProxyRoutingHandler { + private readonly direct = new NodeHttpHandler() + private readonly proxied: NodeHttpHandler + + constructor(proxyUrl: string) { + this.proxied = new NodeHttpHandler({ + httpAgent: new HttpProxyAgent(proxyUrl), + httpsAgent: new HttpsProxyAgent(proxyUrl), + }) + } + + private handlerFor(request: HandleParameters[0]): NodeHttpHandler { + // NO_PROXY entries match on host alone, so the port is left out. + return isNoProxyHost(`${request.protocol}//${request.hostname}`) ? this.direct : this.proxied + } + + handle(...args: HandleParameters) { + return this.handlerFor(args[0]).handle(...args) + } + + updateHttpClientConfig(...args: UpdateClientConfigParameters) { + this.direct.updateHttpClientConfig(...args) + this.proxied.updateHttpClientConfig(...args) + } + + httpHandlerConfigs() { + return this.proxied.httpHandlerConfigs() + } + + destroy() { + this.direct.destroy() + this.proxied.destroy() + } +} + +/** + * Build a request handler routing AWS SDK traffic through the system proxy, except for the + * destinations NO_PROXY excludes. + * + * Returns undefined when no proxy is configured, so the client keeps its own default handler. + */ +export function createProxyRoutingRequestHandler(): ProxyRoutingHandler | undefined { + const proxyUrl = getSystemProxyUrl() + return proxyUrl ? new ProxyRoutingRequestHandler(proxyUrl) : undefined +} + /** * Log a message to the output channel if available. */