R2: per-key rate limiting. After successful auth, the * filter checks the sliding-window limiter. Exceeding * {@code rateLimitPerMin} returns 429. + * + *
R5 / R7: SSE token fallback is scope-limited. The + * {@code ?token=} query param is accepted only on SSE stream paths + * ({@link #isSseStreamPath}), where browser EventSource cannot set an + * Authorization header (R7). It is rejected on every other path so the API + * key never leaks into access / proxy logs for normal calls (R5). Application + * logging here uses {@code getRequestURI()} (no query string), so the key does + * not reach app logs — but a reverse proxy may still log the query string, so + * keep the fallback as narrow as possible. */ @Slf4j @Component @@ -48,9 +57,12 @@ protected void doFilterInternal(HttpServletRequest request, return; } - String token = extractBearerToken(request); + boolean sse = isSseStreamPath(request); + String token = extractToken(request, sse); if (!StringUtils.hasText(token)) { - sendUnauthorized(response, "Missing API key"); + sendUnauthorized(response, sse + ? "Missing API key (Authorization header or ?token= for EventSource)" + : "Missing API key"); return; } @@ -62,8 +74,12 @@ protected void doFilterInternal(HttpServletRequest request, KbApiKeyContext context = authResult.get().context(); - // R2: rate limit check - if (!rateLimiter.tryAcquire(context.keyId(), context.rateLimitPerMin(), Instant.now())) { + // R2: rate limit check — but NOT on the SSE stream path. EventSource + // reconnects/heartbeats would otherwise burn the per-minute window and + // can 429 the key's own POST /research start. Rate limiting belongs on + // the cost-producing endpoints (start/status/cancel), not the progress + // subscription. Per-key concurrency is still enforced upstream. + if (!sse && !rateLimiter.tryAcquire(context.keyId(), context.rateLimitPerMin(), Instant.now())) { sendTooManyRequests(response, context.rateLimitPerMin()); return; } @@ -82,14 +98,33 @@ private boolean isOpenApiPath(HttpServletRequest request) { return uri.startsWith("/api/v1/open/kb/"); } - private String extractBearerToken(HttpServletRequest request) { + /** + * SSE progress stream paths — the only place {@code ?token=} is accepted, + * because browser EventSource cannot set an Authorization header (R7). + * Matched on URI suffix + content type so the fallback tracks whichever + * endpoints expose SSE, without hard-coding a single kbId/sessionId. + */ + private boolean isSseStreamPath(HttpServletRequest request) { + String uri = request.getRequestURI(); + return uri.startsWith("/api/v1/open/kb/") && uri.endsWith("/stream"); + } + + /** + * Extract the API key. Header is always accepted; the {@code ?token=} + * query param is accepted only on SSE stream paths ({@code sse}), + * to keep the key out of access/proxy logs on every other request (R5). + */ + private String extractToken(HttpServletRequest request, boolean sse) { String bearer = request.getHeader("Authorization"); if (StringUtils.hasText(bearer) && bearer.startsWith("Bearer ")) { return bearer.substring(7).trim(); } - // TODO: add ?token= SSE fallback once Deep Research SSE endpoint is live. - // EventSource can't set custom headers; for now P0-A has no SSE path so - // query param would leak the key into access / proxy logs (R5). + if (sse) { + String queryToken = request.getParameter("token"); + if (StringUtils.hasText(queryToken)) { + return queryToken.trim(); + } + } return null; } diff --git a/mateclaw-server/src/main/java/vip/mate/kbopen/controller/KbOpenResearchController.java b/mateclaw-server/src/main/java/vip/mate/kbopen/controller/KbOpenResearchController.java new file mode 100644 index 000000000..5be1986b8 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/kbopen/controller/KbOpenResearchController.java @@ -0,0 +1,224 @@ +package vip.mate.kbopen.controller; + +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.tags.Tag; +import jakarta.servlet.http.HttpServletRequest; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.http.MediaType; +import org.springframework.web.bind.annotation.*; +import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; +import vip.mate.channel.web.ChatStreamTracker; +import vip.mate.channel.web.Utf8SseEmitter; +import vip.mate.common.result.R; +import vip.mate.exception.MateClawException; +import vip.mate.kbopen.auth.KbApiKeyContext; +import vip.mate.kbopen.auth.RequireKbScope; +import vip.mate.kbopen.research.KbResearchSessionRegistry; +import vip.mate.kbopen.research.KbResearchSessionRegistry.Session; +import vip.mate.kbopen.research.KbResearchSessionRegistry.Status; +import vip.mate.kbopen.research.KbResearchSessionRegistry.TooManyConcurrentException; +import vip.mate.wiki.service.WikiKnowledgeBaseService; +import vip.mate.wiki.service.WikiResearchService; +import vip.mate.wiki.service.WikiResearchService.ResearchResult; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +/** + * KB Open API — Deep Research endpoints. + * + *
Unlike the synchronous read endpoints, research is async (multi-step LLM + * pipeline) with SSE progress. The start endpoint returns a sessionId; the + * caller subscribes to SSE for progress, or polls status for the final result. + * + *
R7: the SSE endpoint uses {@code ?token=} query param because browser + * EventSource cannot set Authorization headers. The {@code KbOpenApiAuthFilter} + * already falls back to query param tokens. + * + *
Each session records the owning API key id + kbId so that status/cancel + * endpoints can authorize access (a caller can only query/cancel their own + * sessions). Results are stored on completion for the status endpoint to + * return synchronously. + * + *
This is an in-memory registry (single-node). For multi-node, sessions + * would need to live in a shared store — but research is short-lived (< 1 min + * typical) and the SSE stream must connect to the node running the job, so + * sticky routing is a prerequisite anyway. + * + *
Atomic: {@code incrementAndGet} + rollback on overflow, so concurrent
+ * starts for the same key cannot both slip past the cap. The previous
+ * stream-and-count impl had a check-then-act race.
+ *
+ * @throws TooManyConcurrentException if {@code keyId} already has
+ * {@code maxConcurrentPerKey} RUNNING sessions.
+ */
+ public void startIfAllowed(String sessionId, Long keyId, Long kbId, String topic) {
+ AtomicInteger count = runningPerKey.computeIfAbsent(keyId, k -> new AtomicInteger());
+ int now = count.incrementAndGet();
+ if (now > maxConcurrentPerKey) {
+ count.decrementAndGet(); // rollback — slot was not granted
+ throw new TooManyConcurrentException(
+ "API key already has " + maxConcurrentPerKey
+ + " running research session(s); limit is " + maxConcurrentPerKey);
+ }
+ sessions.put(sessionId, Session.running(sessionId, keyId, kbId, topic));
+ }
+
+ public Optional
+ *
+ */
+class KbOpenApiAuthFilterTest {
+
+ private static final String KEY = "mck_abcd1234";
+ private static final KbApiKeyContext CTX =
+ new KbApiKeyContext(7L, 1L, Set.of(10L), Set.of("kb:search"), 60);
+
+ private KbApiKeyService keyService;
+ private KbApiKeyRateLimiter rateLimiter;
+ private KbOpenApiAuthFilter filter;
+
+ @BeforeEach
+ void setUp() {
+ keyService = mock(KbApiKeyService.class);
+ rateLimiter = mock(KbApiKeyRateLimiter.class);
+ filter = new KbOpenApiAuthFilter(keyService, rateLimiter);
+ when(keyService.authenticate(KEY)).thenReturn(Optional.of(new AuthResult(CTX, LocalDateTime.now())));
+ when(rateLimiter.tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any())).thenReturn(true);
+ }
+
+ private MockHttpServletRequest startRequest() {
+ MockHttpServletRequest req = new MockHttpServletRequest();
+ req.setRequestURI("/api/v1/open/kb/10/research");
+ req.setMethod("POST");
+ req.addHeader("Authorization", "Bearer " + KEY);
+ return req;
+ }
+
+ private MockHttpServletRequest sseRequest() {
+ MockHttpServletRequest req = new MockHttpServletRequest();
+ req.setRequestURI("/api/v1/open/kb/10/research/open-research-x/stream");
+ req.setMethod("GET");
+ req.setQueryString("token=" + KEY);
+ req.addParameter("token", KEY);
+ return req;
+ }
+
+ private int run(MockHttpServletRequest req) throws Exception {
+ MockHttpServletResponse res = new MockHttpServletResponse();
+ filter.doFilter(req, res, new MockFilterChain());
+ return res.getStatus();
+ }
+
+ @Test
+ @DisplayName("non-SSE path: header auth passes")
+ void nonSseHeaderAuth() throws Exception {
+ assertThat(run(startRequest())).isEqualTo(200);
+ }
+
+ @Test
+ @DisplayName("non-SSE path: ?token= is rejected even with a valid key (R5 — no log leak)")
+ void nonSseQueryTokenRejected() throws Exception {
+ MockHttpServletRequest req = startRequest();
+ req.removeHeader("Authorization");
+ req.addParameter("token", KEY);
+ req.setQueryString("token=" + KEY);
+
+ assertThat(run(req)).isEqualTo(401);
+ verify(keyService, never()).authenticate(KEY);
+ }
+
+ @Test
+ @DisplayName("SSE path: ?token= authenticates (R7 — EventSource fallback)")
+ void sseQueryTokenAccepted() throws Exception {
+ assertThat(run(sseRequest())).isEqualTo(200);
+ verify(keyService).authenticate(KEY);
+ }
+
+ @Test
+ @DisplayName("SSE path: missing token → 401")
+ void sseMissingToken() throws Exception {
+ MockHttpServletRequest req = sseRequest();
+ req.removeParameter("token");
+ req.setQueryString(null);
+ assertThat(run(req)).isEqualTo(401);
+ }
+
+ @Test
+ @DisplayName("SSE path: bypasses the per-minute rate limiter (reconnects must not burn quota)")
+ void sseSkipsRateLimit() throws Exception {
+ run(sseRequest());
+ verify(rateLimiter, never())
+ .tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
+ }
+
+ @Test
+ @DisplayName("non-SSE path: still goes through the rate limiter")
+ void nonSseHitsRateLimit() throws Exception {
+ run(startRequest());
+ verify(rateLimiter)
+ .tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
+ }
+}
diff --git a/mateclaw-server/src/test/java/vip/mate/kbopen/research/KbResearchSessionRegistryTest.java b/mateclaw-server/src/test/java/vip/mate/kbopen/research/KbResearchSessionRegistryTest.java
new file mode 100644
index 000000000..fd31b0363
--- /dev/null
+++ b/mateclaw-server/src/test/java/vip/mate/kbopen/research/KbResearchSessionRegistryTest.java
@@ -0,0 +1,256 @@
+package vip.mate.kbopen.research;
+
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import vip.mate.kbopen.research.KbResearchSessionRegistry.Session;
+import vip.mate.kbopen.research.KbResearchSessionRegistry.Status;
+import vip.mate.kbopen.research.KbResearchSessionRegistry.TooManyConcurrentException;
+import vip.mate.wiki.service.WikiResearchService.ResearchResult;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests for {@link KbResearchSessionRegistry} — session lifecycle, status
+ * transitions, sticky CANCELLED terminal, per-key concurrency cap, and TTL
+ * eviction.
+ */
+class KbResearchSessionRegistryTest {
+
+ private static final ResearchResult RESULT = new ResearchResult("topic", List.of(), "final report");
+
+ private KbResearchSessionRegistry newRegistry() {
+ return new KbResearchSessionRegistry(3, Duration.ofMinutes(30));
+ }
+
+ @Test
+ @DisplayName("startIfAllowed creates a RUNNING session")
+ void registerCreatesRunning() {
+ KbResearchSessionRegistry registry = newRegistry();
+ registry.startIfAllowed("s1", 100L, 10L, "test topic");
+
+ Optional