From 2e3dd071a56edcd8e13ca60642d307cc37f86f8f Mon Sep 17 00:00:00 2001
From: matevip <matevip@mate.vip>
Date: Thu, 2 Jul 2026 17:51:25 +0800
Subject: [PATCH] fix(kb-open): assert session belongs to path kbId + cleanup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- requireSessionOwnership now also checks session.kbId() == path kbId (404 on
  mismatch), so a research session started under one KB cannot be addressed via
  another KB path even when the caller's key is bound to both — defense-in-depth
  on top of the keyId ownership check.
- Drop internal "R7" / "review #446" markers from the controller Javadoc in
  favour of functional wording.
- Import Set/Map/concurrent types and static any() instead of inline FQNs in the
  new kb-open research/auth tests, per code style.
---
 .../controller/KbOpenResearchController.java   | 18 ++++++++++++------
 .../kbopen/auth/KbOpenApiAuthFilterTest.java   |  7 ++++---
 .../KbResearchSessionRegistryTest.java         | 11 +++++++----
 3 files changed, 23 insertions(+), 13 deletions(-)

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
index 5be1986b..89565fe5 100644
--- a/mateclaw-server/src/main/java/vip/mate/kbopen/controller/KbOpenResearchController.java
+++ b/mateclaw-server/src/main/java/vip/mate/kbopen/controller/KbOpenResearchController.java
@@ -36,11 +36,11 @@ import java.util.concurrent.Executors;
  * pipeline) with SSE progress. The start endpoint returns a sessionId; the
  * caller subscribes to SSE for progress, or polls status for the final result.
  *
- * <p>R7: the SSE endpoint uses {@code ?token=} query param because browser
+ * <p>The SSE endpoint uses a {@code ?token=} query param because browser
  * EventSource cannot set Authorization headers. The {@code KbOpenApiAuthFilter}
  * already falls back to query param tokens.
  *
- * <h3>Cost &amp; lifecycle controls (review #446)</h3>
+ * <h3>Cost &amp; lifecycle controls</h3>
  * <ul>
  *   <li>Cancel is <b>cooperative</b>: it calls {@link ChatStreamTracker#requestStop}
  *       so {@link WikiResearchService} bails at the next stage boundary — the
@@ -128,7 +128,7 @@ public class KbOpenResearchController {
             @PathVariable Long kbId,
             @PathVariable String sessionId,
             HttpServletRequest request) {
-        requireSessionOwnership(request, sessionId);
+        requireSessionOwnership(request, kbId, sessionId);
 
         SseEmitter emitter = new Utf8SseEmitter(10 * 60 * 1000L);
         boolean attached = streamTracker.attach(sessionId, emitter);
@@ -154,7 +154,7 @@ public class KbOpenResearchController {
             @PathVariable Long kbId,
             @PathVariable String sessionId,
             HttpServletRequest request) {
-        Session session = requireSessionOwnership(request, sessionId);
+        Session session = requireSessionOwnership(request, kbId, sessionId);
         Status status = session.status();
         ResearchResult result = session.result();
 
@@ -181,7 +181,7 @@ public class KbOpenResearchController {
             @PathVariable Long kbId,
             @PathVariable String sessionId,
             HttpServletRequest request) {
-        Session session = requireSessionOwnership(request, sessionId);
+        Session session = requireSessionOwnership(request, kbId, sessionId);
         if (session.status() != Status.RUNNING) {
             throw new MateClawException(409, "Session is not running (status: " + session.status() + ")");
         }
@@ -209,7 +209,7 @@ public class KbOpenResearchController {
         return ctx;
     }
 
-    private Session requireSessionOwnership(HttpServletRequest request, String sessionId) {
+    private Session requireSessionOwnership(HttpServletRequest request, Long kbId, String sessionId) {
         KbApiKeyContext ctx = requireContext(request);
         Optional<Session> session = sessionRegistry.get(sessionId);
         if (session.isEmpty()) {
@@ -219,6 +219,12 @@ public class KbOpenResearchController {
         if (!session.get().keyId().equals(ctx.keyId())) {
             throw new MateClawException(403, "Session does not belong to this API key");
         }
+        // The session must also belong to the KB named in the path, so a session
+        // started under one KB cannot be addressed via another — even when the
+        // caller's key happens to be bound to both.
+        if (!session.get().kbId().equals(kbId)) {
+            throw new MateClawException(404, "Research session not found: " + sessionId);
+        }
         return session.get();
     }
 }
diff --git a/mateclaw-server/src/test/java/vip/mate/kbopen/auth/KbOpenApiAuthFilterTest.java b/mateclaw-server/src/test/java/vip/mate/kbopen/auth/KbOpenApiAuthFilterTest.java
index 0b47d8c2..c9eb94a6 100644
--- a/mateclaw-server/src/test/java/vip/mate/kbopen/auth/KbOpenApiAuthFilterTest.java
+++ b/mateclaw-server/src/test/java/vip/mate/kbopen/auth/KbOpenApiAuthFilterTest.java
@@ -13,6 +13,7 @@ import java.util.Optional;
 import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.mock;
@@ -48,7 +49,7 @@ class KbOpenApiAuthFilterTest {
         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);
+        when(rateLimiter.tryAcquire(anyLong(), anyInt(), any())).thenReturn(true);
     }
 
     private MockHttpServletRequest startRequest() {
@@ -113,7 +114,7 @@ class KbOpenApiAuthFilterTest {
     void sseSkipsRateLimit() throws Exception {
         run(sseRequest());
         verify(rateLimiter, never())
-                .tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
+                .tryAcquire(anyLong(), anyInt(), any());
     }
 
     @Test
@@ -121,6 +122,6 @@ class KbOpenApiAuthFilterTest {
     void nonSseHitsRateLimit() throws Exception {
         run(startRequest());
         verify(rateLimiter)
-                .tryAcquire(anyLong(), anyInt(), org.mockito.ArgumentMatchers.any());
+                .tryAcquire(anyLong(), anyInt(), 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
index fd31b036..75f23c20 100644
--- a/mateclaw-server/src/test/java/vip/mate/kbopen/research/KbResearchSessionRegistryTest.java
+++ b/mateclaw-server/src/test/java/vip/mate/kbopen/research/KbResearchSessionRegistryTest.java
@@ -9,8 +9,11 @@ import vip.mate.wiki.service.WikiResearchService.ResearchResult;
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -187,10 +190,10 @@ class KbResearchSessionRegistryTest {
         int cap = 3;
         KbResearchSessionRegistry registry = new KbResearchSessionRegistry(cap, Duration.ofMinutes(30));
         int threads = cap * 4; // far more contenders than slots
-        java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1);
-        java.util.concurrent.atomic.AtomicInteger admitted = new java.util.concurrent.atomic.AtomicInteger();
-        java.util.concurrent.atomic.AtomicInteger rejected = new java.util.concurrent.atomic.AtomicInteger();
-        java.util.List<Thread> workers = new java.util.ArrayList<>();
+        CountDownLatch start = new CountDownLatch(1);
+        AtomicInteger admitted = new AtomicInteger();
+        AtomicInteger rejected = new AtomicInteger();
+        List<Thread> workers = new ArrayList<>();
 
         for (int i = 0; i < threads; i++) {
             String sid = "concurrent-" + i;