Fix pagination and propagate authentication context in requests

app-builder/plugins/data-mate-knowledge/pom.xml

@@ -139,4 +139,4 @@
             </plugin>
         </plugins>
     </build>
-</project>
+</project>

app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/dto/DataMateKnowledgeListQueryParam.java

@@ -21,7 +21,7 @@
 @SerializeStrategy(include = SerializeStrategy.Include.NON_NULL)
 public class DataMateKnowledgeListQueryParam {
     /**
-     * 页码，从0开始。
+     * 页码，从1开始。
      */
     private Integer page;
 
@@ -40,4 +40,3 @@ public class DataMateKnowledgeListQueryParam {
      */
     private String description;
 }
-

app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/external/DataMateKnowledgeBaseManager.java

@@ -35,6 +35,8 @@
 import modelengine.fitframework.util.MapBuilder;
 import modelengine.fitframework.util.ObjectUtils;
 import modelengine.fitframework.util.StringUtils;
+import modelengine.jade.authentication.context.UserContext;
+import modelengine.jade.authentication.context.UserContextHolder;
 import modelengine.jade.knowledge.code.KnowledgeManagerRetCode;
 import modelengine.jade.knowledge.exception.KnowledgeException;
 
@@ -52,6 +54,8 @@ public class DataMateKnowledgeBaseManager {
     private static final Logger log = Logger.get(DataMateKnowledgeBaseManager.class);
     private static final String BEARER = "Bearer ";
     private static final String CONTENT_TYPE_JSON = "application/json";
+    private static final String USER_HEADER = "User";
+    private static final String RETRIEVE_USER = "admin";
     /** 默认访问超时时间（秒）。 */
     private static final int DEFAULT_TIMEOUT_SECONDS = 30;
 
@@ -89,6 +93,7 @@ public DataMateKnowledgeListEntity listRepos(String apiKey, DataMateKnowledgeLis
         if (StringUtils.isNotEmpty(apiKey)) {
             request.headers().set(AUTHORIZATION, BEARER + apiKey);
         }
+        this.addCurrentUserHeader(request);
         try {
             Object object = this.httpClient.get().exchangeForEntity(request, Object.class);
             Map<String, Object> response =
@@ -118,6 +123,7 @@ public DataMateRetrievalResult retrieve(String apiKey, DataMateRetrievalParam pa
         if (StringUtils.isNotEmpty(apiKey)) {
             request.headers().set(AUTHORIZATION, BEARER + apiKey);
         }
+        this.addUserHeader(request, RETRIEVE_USER);
         request.headers().set(CONTENT_TYPE, CONTENT_TYPE_JSON);
         try {
             Object object = this.httpClient.get().exchangeForEntity(request, Object.class);
@@ -141,6 +147,20 @@ private KnowledgeException handleException(HttpClientResponseException ex) {
         return new KnowledgeException(retCode, ex, ex.getSimpleMessage());
     }
 
+    private void addCurrentUserHeader(HttpClassicClientRequest request) {
+        UserContext context = UserContextHolder.get();
+        String userName = context == null ? StringUtils.EMPTY : context.getName();
+        this.addUserHeader(request, userName);
+    }
+
+    private void addUserHeader(HttpClassicClientRequest request, String userName) {
+        log.info("DataMate request user: {}", StringUtils.isEmpty(userName) ? "<empty>" : userName);
+        if (StringUtils.isEmpty(userName)) {
+            return;
+        }
+        request.headers().set(USER_HEADER, userName);
+    }
+
     private HttpClassicClient getHttpClient() {
         int timeoutMs = this.timeoutSeconds * 1000;
         Map<String, Object> custom = MapBuilder.<String, Object>get()
@@ -155,4 +175,3 @@ private HttpClassicClient getHttpClient() {
                 .build());
     }
 }
-

app-builder/plugins/data-mate-knowledge/src/main/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImpl.java

@@ -132,7 +132,7 @@ public List<KnowledgeDocument> retrieve(String apiKey, FlatKnowledgeOption optio
     }
 
     private PageVoKnowledgeList queryKnowledgeList(String apiKey, ListRepoQueryParam param) {
-        int page = Math.max(param.getPageIndex() - 1, 0);
+        int page = Math.max(param.getPageIndex(), 1);
         int size = param.getPageSize();
         DataMateKnowledgeListEntity listEntity = this.executeQuery(apiKey, param.getRepoName(), page, size);
         List<DataMateKnowledgeEntity> content = listEntity.getContent() == null

app-builder/plugins/data-mate-knowledge/src/main/resources/application.yml

@@ -6,7 +6,6 @@ fit:
 # timeout: 访问超时时间（秒），默认 30，可按场景自定义
 datamate:
   url:
-    list: 'http://datamate-gateway:8080/api/knowledge-base/list'
-    retrieve: 'http://datamate-gateway:8080/api/knowledge-base/retrieve'
+    list: 'http://datamate-backend-python:18000/api/knowledge-base/list'
+    retrieve: 'http://datamate-backend-python:18000/api/knowledge-base/retrieve'
   # timeout: 30
-

app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/DataMateKnowledgeBaseManagerTest.java

@@ -17,6 +17,8 @@
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalChunksEntity;
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalResult;
 import modelengine.fit.jade.datamate.knowledge.external.DataMateKnowledgeBaseManager;
+import modelengine.jade.authentication.context.UserContext;
+import modelengine.jade.authentication.context.UserContextHolder;
 import modelengine.jade.knowledge.exception.KnowledgeException;
 
 import org.junit.jupiter.api.AfterEach;
@@ -31,6 +33,7 @@
 import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * 表示 {@link DataMateKnowledgeBaseManager} 的测试集。
@@ -81,6 +84,28 @@ public void shouldOkWhenListRepo() {
                 DataMateKnowledgeEntity::getDescription).containsExactly("2", "test2", "test2知识库");
     }
 
+    @Test
+    @DisplayName("查询知识库列表时携带 User 请求头")
+    public void shouldAddUserHeaderWhenListRepo() {
+        DataMateKnowledgeListQueryParam param = DataMateKnowledgeListQueryParam.builder().name("user").build();
+        DataMateKnowledgeListEntity entity = this.executeListWithContext(param,
+                new UserContext("test-user", "", ""));
+
+        assertThat(entity.getContent()).hasSize(2);
+        assertThat(UserContextHolder.get()).isNull();
+    }
+
+    @Test
+    @DisplayName("用户名缺失时不添加 User 请求头")
+    public void shouldOmitUserHeaderWhenUserNameMissing() {
+        DataMateKnowledgeListQueryParam param =
+                DataMateKnowledgeListQueryParam.builder().name("missing-user").build();
+        DataMateKnowledgeListEntity entity = this.executeListWithContext(param,
+                new UserContext("", "", ""));
+
+        assertThat(entity.getContent()).hasSize(2);
+    }
+
     @Test
     @DisplayName("查询知识库列表失败，抛出异常")
     public void shouldFailWhenListRepoThrowException() {
@@ -100,6 +125,19 @@ public void shouldOkWhenRetrieve() {
         assertThat(result.getData().get(0).content()).isEqualTo("content1");
     }
 
+    @Test
+    @DisplayName("检索知识库时固定携带 admin 用户")
+    public void shouldAddAdminUserHeaderWhenRetrieve() {
+        DataMateRetrievalParam param = DataMateRetrievalParam.builder().query("admin").build();
+        AtomicReference<DataMateRetrievalResult> result = new AtomicReference<>();
+        UserContext context = new UserContext("test-user", "", "");
+
+        UserContextHolder.apply(context, () -> result.set(this.manager.retrieve(this.apiKey, param)));
+
+        assertThat(result.get().getData()).hasSize(3);
+        assertThat(UserContextHolder.get()).isNull();
+    }
+
     @Test
     @DisplayName("检索知识库失败，抛出异常")
     public void shouldFailWhenRetrieveThrowException() {
@@ -108,4 +146,11 @@ public void shouldFailWhenRetrieveThrowException() {
                 .extracting("code")
                 .isEqualTo(130703005);
     }
+
+    private DataMateKnowledgeListEntity executeListWithContext(DataMateKnowledgeListQueryParam param,
+            UserContext context) {
+        AtomicReference<DataMateKnowledgeListEntity> entity = new AtomicReference<>();
+        UserContextHolder.apply(context, () -> entity.set(this.manager.listRepos(this.apiKey, param)));
+        return entity.get();
+    }
 }

app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeBaseInnerController.java

@@ -6,12 +6,14 @@
 
 package modelengine.fit.jade.datamate.knowledge;
 
+import static modelengine.fit.http.protocol.MessageHeaderNames.AUTHORIZATION;
 import static modelengine.fitframework.util.IoUtils.content;
 
 import modelengine.fit.http.annotation.PostMapping;
 import modelengine.fit.http.annotation.RequestBody;
 import modelengine.fit.http.annotation.RequestMapping;
 import modelengine.fit.http.client.HttpClientException;
+import modelengine.fit.http.server.HttpClassicServerRequest;
 import modelengine.fitframework.annotation.Component;
 import modelengine.fitframework.serialization.ObjectSerializer;
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateResponse;
@@ -28,29 +30,55 @@
 @Component
 @RequestMapping(path = "/v2", group = "DataMate知识库内部接口打桩")
 public class MockedDataMateKnowledgeBaseInnerController {
+    private static final String USER_HEADER = "User";
+    private static final String EXPECTED_AUTHORIZATION = "Bearer 123";
+    private static final String EXPECTED_USER = "test-user";
+    private static final String EXPECTED_RETRIEVE_USER = "admin";
+
     private final ObjectSerializer serializer;
 
     public MockedDataMateKnowledgeBaseInnerController(ObjectSerializer serializer) {
         this.serializer = serializer;
     }
 
     @PostMapping(path = "/knowledgeBase")
-    public Map<String, Object> listRepos(@RequestBody MockedDataMateKnowledgeListQueryParam param) throws IOException {
-        if (param.getName().equals("error")) {
+    public Map<String, Object> listRepos(HttpClassicServerRequest request,
+            @RequestBody MockedDataMateKnowledgeListQueryParam param) throws IOException {
+        if ("error".equals(param.getName())) {
             throw new HttpClientException("error");
         }
+        this.validateRequestContext(request, "user".equals(param.getName()) ? EXPECTED_USER : null);
         String resourceName = "/listRepoResult.json";
         String jsonContent = content(DataMateResponse.class, resourceName);
         return serializer.deserialize(jsonContent, Map.class);
     }
 
     @PostMapping(path = "/knowledgebases/query")
-    public Map<String, Object> retrieve(@RequestBody MockedDataMateRetrievalParam param) throws IOException {
-        if (param.getQuery().equals("error")) {
+    public Map<String, Object> retrieve(HttpClassicServerRequest request,
+            @RequestBody MockedDataMateRetrievalParam param) throws IOException {
+        if ("error".equals(param.getQuery())) {
             throw new HttpClientException("error");
         }
+        this.validateRequestContext(request, EXPECTED_RETRIEVE_USER);
         String resourceName = "/retrieveResult.json";
         String jsonContent = content(DataMateResponse.class, resourceName);
         return serializer.deserialize(jsonContent, Map.class);
     }
+
+    private void validateRequestContext(HttpClassicServerRequest request, String expectedUser) {
+        this.validate(EXPECTED_AUTHORIZATION.equals(request.headers().first(AUTHORIZATION).orElse(null)),
+                "The authorization header is incorrect.");
+        if (expectedUser != null) {
+            this.validate(expectedUser.equals(request.headers().first(USER_HEADER).orElse(null)),
+                    "The User header is incorrect.");
+            return;
+        }
+        this.validate(request.headers().first(USER_HEADER).isEmpty(), "The User header should be absent.");
+    }
+
+    private void validate(boolean expression, String message) {
+        if (!expression) {
+            throw new HttpClientException(message);
+        }
+    }
 }

app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/MockedDataMateKnowledgeListQueryParam.java

@@ -19,7 +19,7 @@
 @SerializeStrategy(include = SerializeStrategy.Include.NON_NULL)
 public class MockedDataMateKnowledgeListQueryParam {
     /**
-     * 页码，从0开始。
+     * 页码，从1开始。
      */
     private Integer page;
 

app-builder/plugins/data-mate-knowledge/src/test/java/modelengine/fit/jade/datamate/knowledge/service/DataMateKnowledgeRepoServiceImplTest.java

@@ -10,9 +10,12 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import modelengine.fit.jade.datamate.knowledge.dto.DataMateKnowledgeListQueryParam;
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateKnowledgeEntity;
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateKnowledgeListEntity;
 import modelengine.fit.jade.datamate.knowledge.entity.DataMateRetrievalChunksEntity;
@@ -36,6 +39,7 @@
 
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 
 import java.util.Collections;
 import java.util.List;
@@ -70,7 +74,7 @@ void shouldReturnPageWhenListRepos() {
                 .build();
 
         DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder()
-                .page(0)
+                .page(1)
                 .size(10)
                 .totalElements(1)
                 .totalPages(1)
@@ -83,6 +87,7 @@ void shouldReturnPageWhenListRepos() {
         param.setRepoName("test");
 
         when(knowledgeBaseManager.listRepos(anyString(), any())).thenReturn(listEntity);
+        clearInvocations(knowledgeBaseManager);
 
         PageVo<KnowledgeRepo> result = knowledgeRepoService.listRepos(API_KEY, param);
 
@@ -91,13 +96,17 @@ void shouldReturnPageWhenListRepos() {
         assertThat(result.getItems().get(0))
                 .extracting(KnowledgeRepo::id, KnowledgeRepo::name, KnowledgeRepo::description, KnowledgeRepo::type)
                 .containsExactly("kb-1", "测试知识库", "描述", "embedding-v1");
+        ArgumentCaptor<DataMateKnowledgeListQueryParam> queryParamCaptor =
+                ArgumentCaptor.forClass(DataMateKnowledgeListQueryParam.class);
+        verify(knowledgeBaseManager).listRepos(anyString(), queryParamCaptor.capture());
+        assertThat(queryParamCaptor.getValue().getPage()).isEqualTo(1);
     }
 
     @Test
     @DisplayName("listRepos - 列表为空时返回空分页")
     void shouldReturnEmptyPageWhenContentEmpty() {
         DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder()
-                .page(0)
+                .page(1)
                 .size(10)
                 .totalElements(0)
                 .totalPages(0)
@@ -120,7 +129,7 @@ void shouldReturnEmptyPageWhenContentEmpty() {
     @DisplayName("listRepos - content 为 null 时按空列表处理")
     void shouldReturnEmptyPageWhenContentNull() {
         DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder()
-                .page(0)
+                .page(1)
                 .size(10)
                 .totalElements(0)
                 .totalPages(0)
@@ -143,7 +152,7 @@ void shouldReturnEmptyPageWhenContentNull() {
     @DisplayName("listRepos - totalElements 为 null 时 total 为 0")
     void shouldUseZeroTotalWhenTotalElementsNull() {
         DataMateKnowledgeListEntity listEntity = DataMateKnowledgeListEntity.builder()
-                .page(0)
+                .page(1)
                 .size(10)
                 .totalElements(null)
                 .totalPages(null)