Skip to content

Commit 9e3d934

Browse files
ppiliaevppiliaev
authored
AmazonQ: Netty client (#4150)
AmazonQ client was changed to netty client to add proxy settings support
1 parent ac76894 commit 9e3d934

File tree

7 files changed

+171
-11
lines changed

7 files changed

+171
-11
lines changed

‎.changes/next-release/bugfix-e737d87c-f84a-44f0-b403-5c692051e9d3.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Respect IDE HTTP proxy server settings when using Amazon Q"
4+
}

‎buildSrc/src/main/kotlin/toolkit-intellij-subplugin.gradle.kts‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,14 +50,6 @@ configurations {
5050
exclude(group = "org.slf4j")
5151
exclude(group = "org.jetbrains.kotlin")
5252
exclude(group = "org.jetbrains.kotlinx")
53-
54-
// Exclude dependencies we don't use to make plugin smaller
55-
exclude(group = "software.amazon.awssdk", module = "netty-nio-client")
56-
}
57-
58-
testRuntimeClasspath {
59-
// Conflicts with CRT in test classpath
60-
exclude(group = "software.amazon.awssdk", module = "netty-nio-client")
6153
}
6254

6355
all {

‎gradle/libs.versions.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ aws-ecs = { module = "software.amazon.awssdk:ecs", version.ref = "awsSdk" }
5050
aws-iam = { module = "software.amazon.awssdk:iam", version.ref = "awsSdk" }
5151
aws-jsonProtocol = { module = "software.amazon.awssdk:aws-json-protocol", version.ref = "awsSdk" }
5252
aws-lambda = { module = "software.amazon.awssdk:lambda", version.ref = "awsSdk" }
53+
aws-nettyClient = { module = "software.amazon.awssdk:netty-nio-client", version.ref = "awsSdk" }
5354
aws-queryProtocol = { module = "software.amazon.awssdk:aws-query-protocol", version.ref = "awsSdk" }
5455
aws-rds = { module = "software.amazon.awssdk:rds", version.ref = "awsSdk" }
5556
aws-redshift = { module = "software.amazon.awssdk:redshift", version.ref = "awsSdk" }

‎plugins/toolkit/intellij/build.gradle.kts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,6 @@ configurations {
9393
exclude(group = "org.slf4j")
9494
exclude(group = "org.jetbrains.kotlin")
9595
exclude(group = "org.jetbrains.kotlinx")
96-
exclude(group = "software.amazon.awssdk", module = "netty-nio-client")
9796
}
9897
}
9998

‎plugins/toolkit/jetbrains-core/build.gradle.kts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ dependencies {
132132
api(libs.aws.ecs)
133133
api(libs.aws.iam)
134134
api(libs.aws.lambda)
135+
api(libs.aws.nettyClient)
135136
api(libs.aws.rds)
136137
api(libs.aws.redshift)
137138
api(libs.aws.s3)
@@ -148,7 +149,6 @@ dependencies {
148149
testImplementation(project(":plugin-core:sdk-codegen"))
149150

150151
implementation(project(":plugin-amazonq:mynah-ui"))
151-
implementation(libs.aws.crt)
152152
implementation(libs.bundles.jackson)
153153
implementation(libs.zjsonpatch)
154154
implementation(libs.commonmark)

‎plugins/toolkit/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererEndpointCustomizer.kt‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
package software.aws.toolkits.jetbrains.services.codewhisperer.util
55

66
import com.intellij.openapi.application.ApplicationManager
7+
import com.intellij.util.net.HttpConfigurable
8+
import com.intellij.util.net.ssl.CertificateManager
9+
import com.intellij.util.proxy.CommonProxy
710
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider
811
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
912
import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider
@@ -14,6 +17,8 @@ import software.amazon.awssdk.core.interceptor.ExecutionAttributes
1417
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor
1518
import software.amazon.awssdk.core.retry.RetryPolicy
1619
import software.amazon.awssdk.http.SdkHttpRequest
20+
import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient
21+
import software.amazon.awssdk.http.nio.netty.ProxyConfiguration
1722
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClientBuilder
1823
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClientBuilder
1924
import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClientBuilder
@@ -23,7 +28,9 @@ import software.aws.toolkits.jetbrains.core.AwsSdkClient
2328
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
2429
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
2530
import software.aws.toolkits.jetbrains.services.telemetry.AwsCognitoCredentialsProvider
31+
import java.net.Proxy
2632
import java.net.URI
33+
import javax.net.ssl.TrustManager
2734

2835
// TODO: move this file to package /client
2936
class CodeWhispererEndpointCustomizer : ToolkitClientCustomizer {
@@ -36,8 +43,9 @@ class CodeWhispererEndpointCustomizer : ToolkitClientCustomizer {
3643
clientOverrideConfiguration: ClientOverrideConfiguration.Builder
3744
) {
3845
if (builder is CodeWhispererRuntimeClientBuilder || builder is CodeWhispererStreamingAsyncClientBuilder) {
46+
val endpoint = URI.create(CodeWhispererConstants.Config.CODEWHISPERER_ENDPOINT)
3947
builder
40-
.endpointOverride(URI.create(CodeWhispererConstants.Config.CODEWHISPERER_ENDPOINT))
48+
.endpointOverride(endpoint)
4149
.region(CodeWhispererConstants.Config.BearerClientRegion)
4250
clientOverrideConfiguration.retryPolicy(RetryPolicy.none())
4351
clientOverrideConfiguration.addExecutionInterceptor(
@@ -56,6 +64,41 @@ class CodeWhispererEndpointCustomizer : ToolkitClientCustomizer {
5664
}
5765
}
5866
)
67+
if (builder is CodeWhispererStreamingAsyncClientBuilder) {
68+
val proxy = CommonProxy.getInstance().select(endpoint).first()
69+
val address = proxy.address()
70+
val clientBuilder = NettyNioAsyncHttpClient.builder()
71+
72+
// proxy.type is one of {DIRECT, HTTP, SOCKS}, and is definitely a InetSocketAddress in the HTTP/SOCKS case
73+
// and is null in DIRECT case
74+
if (proxy.type() == Proxy.Type.SOCKS) {
75+
error("Q Chat HTTP client does not support SOCKS proxies")
76+
} else if (address is java.net.InetSocketAddress) {
77+
val proxyConfiguration = ProxyConfiguration.builder()
78+
.host(address.getHostName())
79+
.port(address.getPort())
80+
.apply {
81+
val configurable = HttpConfigurable.getInstance()
82+
val proxyExceptions = configurable.PROXY_EXCEPTIONS
83+
if (!proxyExceptions.isNullOrBlank()) {
84+
// should be handled by CommonProxy, but also should be no harm in passing this along if something more complicated is happening
85+
nonProxyHosts(proxyExceptions.split(',').toSet())
86+
}
87+
88+
val login = HttpConfigurable.getInstance().proxyLogin
89+
if (login != null) {
90+
username(login)
91+
password(HttpConfigurable.getInstance().plainProxyPassword)
92+
}
93+
}
94+
.useSystemPropertyValues(false)
95+
96+
clientBuilder.proxyConfiguration(proxyConfiguration.build())
97+
.tlsTrustManagersProvider { arrayOf<TrustManager>(CertificateManager.getInstance().trustManager) }
98+
}
99+
100+
builder.httpClientBuilder(clientBuilder)
101+
}
59102
} else if (builder is CodeWhispererClientBuilder) {
60103
clientOverrideConfiguration.addExecutionInterceptor(
61104
object : ExecutionInterceptor {

‎plugins/toolkit/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererEndpointCustomizerTest.kt‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer
5+
6+
import com.intellij.openapi.application.ApplicationManager
7+
import com.intellij.testFramework.ApplicationRule
8+
import com.intellij.testFramework.DisposableRule
9+
import com.intellij.util.net.HttpConfigurable
10+
import org.eclipse.jetty.proxy.ConnectHandler
11+
import org.eclipse.jetty.proxy.ProxyServlet
12+
import org.eclipse.jetty.server.Server
13+
import org.eclipse.jetty.server.ServerConnector
14+
import org.eclipse.jetty.servlet.ServletContextHandler
15+
import org.eclipse.jetty.servlet.ServletHolder
16+
import org.junit.After
17+
import org.junit.Before
18+
import org.junit.Rule
19+
import org.junit.Test
20+
import org.mockito.Mockito.times
21+
import org.mockito.kotlin.any
22+
import org.mockito.kotlin.spy
23+
import org.mockito.kotlin.verify
24+
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider
25+
import software.amazon.awssdk.auth.token.credentials.SdkToken
26+
import software.amazon.awssdk.auth.token.credentials.StaticTokenProvider
27+
import software.amazon.awssdk.regions.Region
28+
import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClient
29+
import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClientBuilder
30+
import software.amazon.awssdk.services.codewhispererstreaming.model.GenerateAssistantResponseResponseHandler
31+
import software.aws.toolkits.jetbrains.core.AwsClientManager
32+
import software.aws.toolkits.jetbrains.core.MockClientManager.Companion.useRealImplementations
33+
import java.util.concurrent.CountDownLatch
34+
35+
class CodeWhispererEndpointCustomizerTest {
36+
@Rule
37+
@JvmField
38+
val application = ApplicationRule()
39+
40+
@Rule
41+
@JvmField
42+
val disposableRule = DisposableRule()
43+
44+
private val proxyServletSpy = spy(ConnectHandler())
45+
46+
private val proxyServer = Server().also {
47+
it.addConnector(ServerConnector(it))
48+
it.handler = proxyServletSpy
49+
50+
val context = ServletContextHandler(proxyServletSpy, "/", ServletContextHandler.SESSIONS)
51+
52+
context.addServlet(ServletHolder(ProxyServlet()), "/*")
53+
it.start()
54+
}
55+
56+
@Before
57+
fun setUp() {
58+
useRealImplementations(disposableRule.disposable)
59+
val httpConfigurable = HttpConfigurable.getInstance()
60+
httpConfigurable.USE_HTTP_PROXY = true
61+
httpConfigurable.PROXY_HOST = "localhost"
62+
httpConfigurable.PROXY_PORT = proxyServer.uri.port
63+
}
64+
65+
@After
66+
fun tearDown() {
67+
HttpConfigurable.getInstance().USE_HTTP_PROXY = false
68+
69+
proxyServer.stop()
70+
proxyServer.join()
71+
}
72+
73+
@Test
74+
fun proxyIsBypassed() {
75+
HttpConfigurable.getInstance().USE_HTTP_PROXY = false
76+
77+
makeAwsCall()
78+
79+
verify(proxyServletSpy, times(0)).handle(any(), any(), any(), any())
80+
}
81+
82+
@Test
83+
fun proxyCallIsMade() {
84+
makeAwsCall()
85+
86+
verify(proxyServletSpy).handle(any(), any(), any(), any())
87+
}
88+
89+
private fun makeAwsCall() {
90+
val latch = CountDownLatch(1)
91+
92+
ApplicationManager.getApplication().executeOnPooledThread {
93+
AwsClientManager.getInstance().createUnmanagedClient<CodeWhispererStreamingAsyncClient>(
94+
AnonymousCredentialsProvider.create(),
95+
Region.AWS_GLOBAL,
96+
clientCustomizer = { _, _, _, builder, _ ->
97+
(builder as CodeWhispererStreamingAsyncClientBuilder).tokenProvider(
98+
StaticTokenProvider.create(
99+
object : SdkToken {
100+
override fun token() = "testToken"
101+
override fun expirationTime() = null
102+
}
103+
)
104+
)
105+
}
106+
)
107+
.use {
108+
it.generateAssistantResponse(
109+
{},
110+
GenerateAssistantResponseResponseHandler.builder()
111+
.onEventStream {}
112+
.onError { latch.countDown() }
113+
.onComplete { latch.countDown() }
114+
.build()
115+
)
116+
}
117+
}
118+
119+
latch.await()
120+
}
121+
}

0 commit comments

Comments
 (0)