Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions src/services/code-index/embedders/__tests__/bedrock.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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: {
Expand Down Expand Up @@ -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")
})
Expand All @@ -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", () => {
Expand Down
7 changes: 7 additions & 0 deletions src/services/code-index/embedders/bedrock.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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 }),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
})

this.defaultModelId = modelId || getDefaultModelId("bedrock")
Expand Down
140 changes: 139 additions & 1 deletion src/utils/__tests__/networkProxy.spec.ts
Original file line number Diff line number Diff line change
@@ -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: {
Expand Down Expand Up @@ -469,4 +493,118 @@ describe("networkProxy", () => {
})
})
})

describe("createProxyRoutingRequestHandler", () => {
type Handler = NonNullable<ReturnType<typeof createProxyRoutingRequestHandler>>

// 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<Handler["handle"]>[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
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This asserts only the first argument. If handle dropped requestOptions (abort signal, requestTimeout), would this test still pass? Calling handler.handle(request, { abortSignal }) and asserting both arguments, plus the returned promise, would pin that down.

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()
})
})
})
76 changes: 76 additions & 0 deletions src/utils/networkProxy.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"

/**
Expand Down Expand Up @@ -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<NodeHttpHandler["handle"]>
type UpdateClientConfigParameters = Parameters<NodeHttpHandler["updateHttpClientConfig"]>

/**
* 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<NodeHttpHandler["handle"]>
updateHttpClientConfig(...args: UpdateClientConfigParameters): void
httpHandlerConfigs(): ReturnType<NodeHttpHandler["httpHandlerConfigs"]>
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),
Comment on lines +458 to +459

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since these agents are built with no options, https-proxy-agent defaults to keepAlive: false, and Smithy only applies its keep-alive defaults to agents it builds itself. Does each proxied embedding request therefore pay a fresh CONNECT and TLS handshake, while the direct handler keeps its connections alive? If that is intended, could the comment say why? Otherwise { keepAlive: true, maxSockets: 50 } on both agents would match the direct route.

})
}

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.
*/
Expand Down
Loading