fix: validate initialize response identity

Author: kamikaz1k

acp-core/src/main/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpAcpClientTransport.java

@@ -282,17 +282,24 @@ private Mono<Void> initialize(AcpSchema.JSONRPCRequest request) {
 					return Mono.error(new AcpConnectionException(
 							"Expected " + CONTENT_TYPE_JSON + " initialize response, got " + contentType));
 				}
-				this.connectionId = response.headers()
-					.firstValue(HEADER_CONNECTION_ID)
-					.orElseThrow(() -> new AcpConnectionException(
-							"Initialize response missing " + HEADER_CONNECTION_ID));
 				JSONRPCMessage responseMessage;
 				try {
 					responseMessage = AcpSchema.deserializeJsonRpcMessage(jsonMapper, response.body());
 				}
 				catch (Exception e) {
 					return Mono.error(new AcpConnectionException("Failed to deserialize initialize response", e));
 				}
+				if (!(responseMessage instanceof AcpSchema.JSONRPCResponse initializeResponse)) {
+					return Mono.error(new AcpConnectionException("ACP initialize response was not a JSON-RPC response"));
+				}
+				if (!Objects.equals(request.id(), initializeResponse.id())) {
+					return Mono.error(
+							new AcpConnectionException("ACP initialize response id did not match initialize request"));
+				}
+				this.connectionId = response.headers()
+					.firstValue(HEADER_CONNECTION_ID)
+					.orElseThrow(() -> new AcpConnectionException(
+							"Initialize response missing " + HEADER_CONNECTION_ID));
 				return openConnectionStream().then(emitInbound(responseMessage));
 			})
 			.doOnError(error -> {

acp-core/src/test/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpAcpClientTransportTest.java

@@ -91,6 +91,35 @@ void defaultAcpPathIsCorrect() {
 		assertThat(StreamableHttpAcpClientTransport.DEFAULT_ACP_PATH).isEqualTo("/acp");
 	}
 
+	@Test
+	void initializeRejectsNonResponseBody() throws Exception {
+		StreamableHttpAcpClientTransport transport = new StreamableHttpAcpClientTransport(
+				URI.create("https://localhost:8443/acp"), jsonMapper,
+				initializeHttpClient(new AcpSchema.JSONRPCNotification("session/update", Map.of())));
+		transport.setExceptionHandler(error -> {
+		});
+
+		assertThatThrownBy(() -> transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_INITIALIZE,
+				"init-1", AcpTestFixtures.createInitializeRequest())).block())
+			.isInstanceOf(AcpConnectionException.class)
+			.hasMessage("ACP initialize response was not a JSON-RPC response");
+	}
+
+	@Test
+	void initializeRejectsResponseWithDifferentId() throws Exception {
+		StreamableHttpAcpClientTransport transport = new StreamableHttpAcpClientTransport(
+				URI.create("https://localhost:8443/acp"), jsonMapper,
+				initializeHttpClient(AcpTestFixtures.createJsonRpcResponse("wrong-id",
+						AcpTestFixtures.createInitializeResponse())));
+		transport.setExceptionHandler(error -> {
+		});
+
+		assertThatThrownBy(() -> transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_INITIALIZE,
+				"init-1", AcpTestFixtures.createInitializeRequest())).block())
+			.isInstanceOf(AcpConnectionException.class)
+			.hasMessage("ACP initialize response id did not match initialize request");
+	}
+
 	@Test
 	void concurrentSessionLoadsReuseInFlightSessionStreamOpen() throws Exception {
 		HttpClient httpClient = mock(HttpClient.class);
@@ -572,6 +601,15 @@ private InputStream emptyBody() {
 		return new ByteArrayInputStream(new byte[0]);
 	}
 
+	private HttpClient initializeHttpClient(AcpSchema.JSONRPCMessage initializeResponse) throws Exception {
+		HttpClient httpClient = mock(HttpClient.class);
+		String body = jsonMapper.writeValueAsString(initializeResponse);
+		HttpResponse<Object> response = response(200,
+				Map.of("Content-Type", "application/json", "Acp-Connection-Id", "conn-1"), body);
+		when(httpClient.sendAsync(any(), any())).thenReturn(CompletableFuture.completedFuture(response));
+		return httpClient;
+	}
+
 	private <T> HttpResponse<T> response(int statusCode, Map<String, String> headers, T body) {
 		HttpResponse<T> response = mock(HttpResponse.class);
 		when(response.statusCode()).thenReturn(statusCode);