mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
feat(kb-open): Deep Research 开放 API（start/SSE/status/cancel） (#446)
* feat(kb-open): Deep Research open API (start/SSE/status/cancel) Implements the async Deep Research endpoint for the KB Open API (#443). Research is a multi-step LLM pipeline (plan → retrieve+draft → compose) that runs asynchronously and broadcasts progress via SSE. Endpoints: - POST /{kbId}/research start (returns sessionId + streamUrl) - GET /{kbId}/research/{id}/stream SSE progress (?token= for EventSource) - GET /{kbId}/research/{id}/status query status / final report - POST /{kbId}/research/{id}/cancel cancel running session Components: - KbOpenResearchController: 4 endpoints, @RequireKbScope("kb:search") - KbResearchSessionRegistry: in-memory session tracking with keyId ownership (a caller can only query/cancel their own sessions) Security: - R7: SSE uses ?token= query param (KbOpenApiAuthFilter already supports this fallback for EventSource which can't set Authorization headers) - Session ownership: status/cancel/stream all verify keyId match - Cancel checks session is RUNNING (409 otherwise) Reuses existing WikiResearchService.research() + ChatStreamTracker for the actual research pipeline and SSE broadcasting. Tests (6 new, all green): - KbResearchSessionRegistryTest: register/complete/fail/cancel lifecycle, cancel-on-completed no-op, unknown session returns empty Closes #443 * fix(kb-open-research): cooperative cancel, sticky terminal, TTL, concurrency cap Review #446 — address all 4 job-lifecycle/cost blockers + nits: 1. Cooperative cancellation (was: cancel only flipped status, pipeline ran to completion). Cancel endpoint now calls streamTracker.requestStop(); WikiResearchService.ensureNotCancelled() checks isStopRequested at each stage boundary (plan→draft, draft→compose) and inside the parallel draft fan-out — so cancel actually halts the expensive LLM calls, not just the SSE stream. Throws ResearchCancelledException (caught locally, no error broadcast). 2. Sticky CANCELLED terminal. complete()/fail() now no-op on a CANCELLED session, so a user who cancelled never sees a COMPLETED report surface via /status. 3. Session registry TTL. Terminal sessions get an updatedAt timestamp and are evicted by a @Scheduled sweep after mate.kbopen.research.session-ttl (default 30m). RUNNING sessions are never evicted. Prevents unbounded memory growth. 4. Per-key concurrency cap. startIfAllowed() rejects new research when a key already has mate.kbopen.research.max-concurrent-per-key (default 3) RUNNING sessions → 429. Stops one key from spawning ~60 parallel multi-step LLM pipelines per minute under the per-min rate limiter. 5. Inline FQN → import (controller LinkedHashMap, test List.of). Nits (inherited from P0-A rebase): - V162→V164, prefix VARCHAR(12), design doc moved to rfcs/. - Design doc: kb:search scope row now documents it covers /research/**. 31 tests pass (12 registry incl. sticky-cancel/concurrency/TTL + 13 service + 4 rate limiter + 4 controller + ...). * fix(kb-open): scope-limited ?token= SSE auth fallback in KbOpenApiAuthFilter R7: the SSE progress stream (/research/{id}/stream) is consumed by browser EventSource, which cannot set an Authorization header. The filter's extractBearerToken() never read ?token= (still a TODO), so the SSE endpoint was unreachable from the browser — the headline use case got 401. Fix: accept ?token= ONLY on SSE stream paths (isSseStreamPath, suffix /stream), reject it everywhere else so the API key does not leak into access/proxy logs for normal calls (R5). Matches the JwtAuthFilter convention (getRequestURI logs carry no query string). Also bypass the per-minute rate limiter on the SSE path: EventSource reconnects/heartbeats would otherwise burn the key's window and 429 its own POST /research start. Rate limiting belongs on the cost-producing endpoints. Tests (6 new, KbOpenApiAuthFilterTest): - non-SSE: header passes, ?token= rejected (no authenticate call) - SSE: ?token= authenticates, missing token → 401 - SSE: bypasses rate limiter; non-SSE still hits it * fix(kb-open-research): make per-key concurrency cap atomic (no check-then-act race) startIfAllowed() did stream-and-count then put() — not atomic. Two concurrent starts for the same key could both pass the count check (both see < cap) and both put, admitting more sessions than the cap. On the virtual-thread start endpoint this is a real DoS/cost-bypass path. Fix: maintain a per-key AtomicInteger running counter (runningPerKey), incremented atomically on start (incrementAndGet + rollback on overflow) and decremented on each RUNNING→terminal transition (complete/fail/cancel). The counter is kept in lock-step with status==RUNNING; since terminal states are sticky, each session decrements exactly once. cancel() also rewritten to capture the pre-transition state cleanly (the old return check relied on Map.computeIfPresent returning the new value, which worked but read as 'before.status==CANCELLED'). Tests (+2): cancelled/failed release slot (counter consistency), and a concurrent-start test (12 virtual threads, cap=3) asserting exactly cap admits — would be flaky/fail under the old impl. * refactor(kb-open-research): remove unused register() back-compat method register() was left over from the initial impl — it bypassed the per-key concurrency cap (no startIfAllowed check) and, after the atomic-counter fix, incremented runningPerKey without any overflow rollback. With no production caller (the start endpoint uses startIfAllowed), it only existed for tests to set up a RUNNING session. Drop it and route the tests through startIfAllowed so nothing can accidentally ship a path that ignores the cap.
This commit is contained in:
倪程伟
GitHub
parent
9d4041714f
commit
20c681a7c8
@ -29,6 +29,15 @@ import java.util.Optional;
|
|||||||
* <p><strong>R2: per-key rate limiting.</strong> After successful auth, the
|
* <p><strong>R2: per-key rate limiting.</strong> After successful auth, the
|
||||||
* filter checks the sliding-window limiter. Exceeding
|
* filter checks the sliding-window limiter. Exceeding
|
||||||
* {@code rateLimitPerMin} returns 429.
|
* {@code rateLimitPerMin} returns 429.
|
||||||
|
*
|
||||||
|
* <p><strong>R5 / R7: SSE token fallback is scope-limited.</strong> The
|
||||||
|
* {@code ?token=} query param is accepted <em>only</em> 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
|
@Slf4j
|
||||||
@Component
|
@Component
|
||||||
@ -48,9 +57,12 @@ public class KbOpenApiAuthFilter extends OncePerRequestFilter {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
String token = extractBearerToken(request);
|
boolean sse = isSseStreamPath(request);
|
||||||
|
String token = extractToken(request, sse);
|
||||||
if (!StringUtils.hasText(token)) {
|
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;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -62,8 +74,12 @@ public class KbOpenApiAuthFilter extends OncePerRequestFilter {
|
|||||||
|
|
||||||
KbApiKeyContext context = authResult.get().context();
|
KbApiKeyContext context = authResult.get().context();
|
||||||
|
|
||||||
// R2: rate limit check
|
// R2: rate limit check — but NOT on the SSE stream path. EventSource
|
||||||
if (!rateLimiter.tryAcquire(context.keyId(), context.rateLimitPerMin(), Instant.now())) {
|
// 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());
|
sendTooManyRequests(response, context.rateLimitPerMin());
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@ -82,14 +98,33 @@ public class KbOpenApiAuthFilter extends OncePerRequestFilter {
|
|||||||
return uri.startsWith("/api/v1/open/kb/");
|
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 <em>only</em> 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");
|
String bearer = request.getHeader("Authorization");
|
||||||
if (StringUtils.hasText(bearer) && bearer.startsWith("Bearer ")) {
|
if (StringUtils.hasText(bearer) && bearer.startsWith("Bearer ")) {
|
||||||
return bearer.substring(7).trim();
|
return bearer.substring(7).trim();
|
||||||
}
|
}
|
||||||
// TODO: add ?token= SSE fallback once Deep Research SSE endpoint is live.
|
if (sse) {
|
||||||
// EventSource can't set custom headers; for now P0-A has no SSE path so
|
String queryToken = request.getParameter("token");
|
||||||
// query param would leak the key into access / proxy logs (R5).
|
if (StringUtils.hasText(queryToken)) {
|
||||||
|
return queryToken.trim();
|
||||||
|
}
|
||||||
|
}
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -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.
|
||||||
|
*
|
||||||
|
* <p>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.
|
||||||
|
*
|
||||||
|
* <p>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.
|
||||||
|
*
|
||||||
|
* <h3>Cost & lifecycle controls (review #446)</h3>
|
||||||
|
* <ul>
|
||||||
|
* <li>Cancel is <b>cooperative</b>: it calls {@link ChatStreamTracker#requestStop}
|
||||||
|
* so {@link WikiResearchService} bails at the next stage boundary — the
|
||||||
|
* draft fan-out and compose LLM calls are skipped, not run to completion.</li>
|
||||||
|
* <li>CANCELLED is a <b>sticky terminal</b>: a late complete() after cancel is
|
||||||
|
* a no-op in the registry, so the user never sees a COMPLETED report.</li>
|
||||||
|
* <li>Per-key <b>concurrency cap</b> ({@code mate.kbopen.research.max-concurrent-per-key})
|
||||||
|
* stops one key from spawning unbounded parallel pipelines → 429.</li>
|
||||||
|
* </ul>
|
||||||
|
*/
|
||||||
|
@Slf4j
|
||||||
|
@Tag(name = "KB Open API — Deep Research")
|
||||||
|
@RestController
|
||||||
|
@RequestMapping("/api/v1/open/kb")
|
||||||
|
@RequiredArgsConstructor
|
||||||
|
public class KbOpenResearchController {
|
||||||
|
|
||||||
|
private final WikiResearchService researchService;
|
||||||
|
private final WikiKnowledgeBaseService kbService;
|
||||||
|
private final ChatStreamTracker streamTracker;
|
||||||
|
private final KbResearchSessionRegistry sessionRegistry;
|
||||||
|
|
||||||
|
private static final ExecutorService RESEARCH_EXEC = Executors.newVirtualThreadPerTaskExecutor();
|
||||||
|
|
||||||
|
// ── POST /{kbId}/research — start ─────────────────────────────────────
|
||||||
|
|
||||||
|
@RequireKbScope("kb:search")
|
||||||
|
@PostMapping("/{kbId}/research")
|
||||||
|
@Operation(summary = "Start Deep Research (async, returns sessionId)")
|
||||||
|
public R<Map<String, Object>> startResearch(
|
||||||
|
@PathVariable Long kbId,
|
||||||
|
@RequestBody ResearchRequest req,
|
||||||
|
HttpServletRequest request) {
|
||||||
|
KbApiKeyContext ctx = requireContext(request);
|
||||||
|
String topic = req.topic();
|
||||||
|
if (topic == null || topic.isBlank()) {
|
||||||
|
throw new MateClawException(400, "topic is required");
|
||||||
|
}
|
||||||
|
if (kbService.getById(kbId) == null) {
|
||||||
|
throw new MateClawException(404, "Knowledge base not found: " + kbId);
|
||||||
|
}
|
||||||
|
|
||||||
|
String sessionId = "open-research-" + UUID.randomUUID();
|
||||||
|
streamTracker.register(sessionId);
|
||||||
|
streamTracker.incrementFlux(sessionId);
|
||||||
|
// Per-key concurrency cap (DoS / runaway-cost guard on top of the
|
||||||
|
// per-minute rate limiter). Throws → 429 below.
|
||||||
|
try {
|
||||||
|
sessionRegistry.startIfAllowed(sessionId, ctx.keyId(), kbId, topic);
|
||||||
|
} catch (TooManyConcurrentException e) {
|
||||||
|
try { streamTracker.complete(sessionId); } catch (Exception ignored) {}
|
||||||
|
throw new MateClawException(429, e.getMessage());
|
||||||
|
}
|
||||||
|
|
||||||
|
RESEARCH_EXEC.submit(() -> {
|
||||||
|
try {
|
||||||
|
ResearchResult result = researchService.research(kbId, topic, sessionId, req.topKPerQuestion());
|
||||||
|
// complete() is a no-op if the user already cancelled — the
|
||||||
|
// sticky CANCELLED terminal wins over a late COMPLETED.
|
||||||
|
sessionRegistry.complete(sessionId, result);
|
||||||
|
} catch (Exception e) {
|
||||||
|
log.error("[KbOpenResearch] Failed sessionId={}: {}", sessionId, e.getMessage(), e);
|
||||||
|
sessionRegistry.fail(sessionId, e.getMessage());
|
||||||
|
} finally {
|
||||||
|
try { streamTracker.broadcast(sessionId, "done", "{}"); } catch (Exception ignored) {}
|
||||||
|
try { streamTracker.complete(sessionId); } catch (Exception ignored) {}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
return R.ok(Map.of(
|
||||||
|
"sessionId", sessionId,
|
||||||
|
"kbId", kbId,
|
||||||
|
"topic", topic,
|
||||||
|
"streamUrl", "/api/v1/open/kb/" + kbId + "/research/" + sessionId + "/stream"));
|
||||||
|
}
|
||||||
|
|
||||||
|
public record ResearchRequest(String topic, Integer topKPerQuestion) {}
|
||||||
|
|
||||||
|
// ── GET /{kbId}/research/{sessionId}/stream — SSE ─────────────────────
|
||||||
|
|
||||||
|
@RequireKbScope("kb:search")
|
||||||
|
@GetMapping(value = "/{kbId}/research/{sessionId}/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
|
||||||
|
@Operation(summary = "Subscribe to research SSE progress (use ?token= for EventSource)")
|
||||||
|
public SseEmitter stream(
|
||||||
|
@PathVariable Long kbId,
|
||||||
|
@PathVariable String sessionId,
|
||||||
|
HttpServletRequest request) {
|
||||||
|
requireSessionOwnership(request, sessionId);
|
||||||
|
|
||||||
|
SseEmitter emitter = new Utf8SseEmitter(10 * 60 * 1000L);
|
||||||
|
boolean attached = streamTracker.attach(sessionId, emitter);
|
||||||
|
if (!attached) {
|
||||||
|
try {
|
||||||
|
emitter.send(SseEmitter.event().name("error")
|
||||||
|
.data("{\"message\":\"session not found or already ended\"}"));
|
||||||
|
emitter.complete();
|
||||||
|
} catch (Exception ignored) {}
|
||||||
|
}
|
||||||
|
emitter.onCompletion(() -> streamTracker.detach(sessionId, emitter));
|
||||||
|
emitter.onTimeout(() -> streamTracker.detach(sessionId, emitter));
|
||||||
|
emitter.onError(err -> streamTracker.detach(sessionId, emitter));
|
||||||
|
return emitter;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── GET /{kbId}/research/{sessionId}/status — query result ────────────
|
||||||
|
|
||||||
|
@RequireKbScope("kb:search")
|
||||||
|
@GetMapping("/{kbId}/research/{sessionId}/status")
|
||||||
|
@Operation(summary = "Query research status / final result")
|
||||||
|
public R<Map<String, Object>> status(
|
||||||
|
@PathVariable Long kbId,
|
||||||
|
@PathVariable String sessionId,
|
||||||
|
HttpServletRequest request) {
|
||||||
|
Session session = requireSessionOwnership(request, sessionId);
|
||||||
|
Status status = session.status();
|
||||||
|
ResearchResult result = session.result();
|
||||||
|
|
||||||
|
Map<String, Object> data = new LinkedHashMap<>();
|
||||||
|
data.put("sessionId", sessionId);
|
||||||
|
data.put("status", status.name().toLowerCase());
|
||||||
|
data.put("topic", session.topic());
|
||||||
|
if (result != null) {
|
||||||
|
data.put("report", result.report());
|
||||||
|
data.put("sections", result.sections().size());
|
||||||
|
}
|
||||||
|
if (session.error() != null) {
|
||||||
|
data.put("error", session.error());
|
||||||
|
}
|
||||||
|
return R.ok(data);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── POST /{kbId}/research/{sessionId}/cancel — cancel ─────────────────
|
||||||
|
|
||||||
|
@RequireKbScope("kb:search")
|
||||||
|
@PostMapping("/{kbId}/research/{sessionId}/cancel")
|
||||||
|
@Operation(summary = "Cancel a running research session")
|
||||||
|
public R<Map<String, Object>> cancel(
|
||||||
|
@PathVariable Long kbId,
|
||||||
|
@PathVariable String sessionId,
|
||||||
|
HttpServletRequest request) {
|
||||||
|
Session session = requireSessionOwnership(request, sessionId);
|
||||||
|
if (session.status() != Status.RUNNING) {
|
||||||
|
throw new MateClawException(409, "Session is not running (status: " + session.status() + ")");
|
||||||
|
}
|
||||||
|
// Cooperative cancellation: signal the running pipeline to bail at the
|
||||||
|
// next stage boundary (plan→draft, draft→compose, and inside the draft
|
||||||
|
// fan-out) rather than running LLM calls to completion.
|
||||||
|
streamTracker.requestStop(sessionId);
|
||||||
|
sessionRegistry.cancel(sessionId);
|
||||||
|
// Close the SSE stream so subscribers detach immediately.
|
||||||
|
try {
|
||||||
|
streamTracker.broadcast(sessionId, "cancelled", "{\"message\":\"cancelled by user\"}");
|
||||||
|
streamTracker.broadcast(sessionId, "done", "{}");
|
||||||
|
streamTracker.complete(sessionId);
|
||||||
|
} catch (Exception ignored) {}
|
||||||
|
return R.ok(Map.of("sessionId", sessionId, "status", "cancelled"));
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Auth helpers ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
private KbApiKeyContext requireContext(HttpServletRequest request) {
|
||||||
|
KbApiKeyContext ctx = (KbApiKeyContext) request.getAttribute(KbApiKeyContext.ATTR);
|
||||||
|
if (ctx == null) {
|
||||||
|
throw new MateClawException(401, "Authentication required");
|
||||||
|
}
|
||||||
|
return ctx;
|
||||||
|
}
|
||||||
|
|
||||||
|
private Session requireSessionOwnership(HttpServletRequest request, String sessionId) {
|
||||||
|
KbApiKeyContext ctx = requireContext(request);
|
||||||
|
Optional<Session> session = sessionRegistry.get(sessionId);
|
||||||
|
if (session.isEmpty()) {
|
||||||
|
throw new MateClawException(404, "Research session not found: " + sessionId);
|
||||||
|
}
|
||||||
|
// A caller can only access sessions they started
|
||||||
|
if (!session.get().keyId().equals(ctx.keyId())) {
|
||||||
|
throw new MateClawException(403, "Session does not belong to this API key");
|
||||||
|
}
|
||||||
|
return session.get();
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,188 @@
|
|||||||
|
package vip.mate.kbopen.research;
|
||||||
|
|
||||||
|
import lombok.extern.slf4j.Slf4j;
|
||||||
|
import org.springframework.beans.factory.annotation.Value;
|
||||||
|
import org.springframework.scheduling.annotation.Scheduled;
|
||||||
|
import org.springframework.stereotype.Component;
|
||||||
|
import vip.mate.wiki.service.WikiResearchService.ResearchResult;
|
||||||
|
|
||||||
|
import java.time.Duration;
|
||||||
|
import java.time.Instant;
|
||||||
|
import java.util.Map;
|
||||||
|
import java.util.Optional;
|
||||||
|
import java.util.concurrent.ConcurrentHashMap;
|
||||||
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Tracks active Deep Research sessions started via the Open API.
|
||||||
|
*
|
||||||
|
* <p>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.
|
||||||
|
*
|
||||||
|
* <p>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.
|
||||||
|
*
|
||||||
|
* <h3>Lifecycle invariants</h3>
|
||||||
|
* <ul>
|
||||||
|
* <li>{@link Status#CANCELLED} is a <b>sticky</b> terminal state: a late
|
||||||
|
* {@link #complete}/{@link #fail} arriving after cancel is a no-op, so
|
||||||
|
* the user who cancelled never sees a COMPLETED report.</li>
|
||||||
|
* <li>Terminal sessions are evicted by {@link #evictExpired} after
|
||||||
|
* {@code mate.kbopen.research.session-ttl} (default 30 min) so the map
|
||||||
|
* cannot grow without bound.</li>
|
||||||
|
* <li>{@link #startIfAllowed} enforces a per-key concurrency cap
|
||||||
|
* ({@code mate.kbopen.research.max-concurrent-per-key}, default 3) as a
|
||||||
|
* DoS / runaway-cost guard on top of the per-minute rate limiter.</li>
|
||||||
|
* </ul>
|
||||||
|
*/
|
||||||
|
@Slf4j
|
||||||
|
@Component
|
||||||
|
public class KbResearchSessionRegistry {
|
||||||
|
|
||||||
|
public enum Status { RUNNING, COMPLETED, FAILED, CANCELLED }
|
||||||
|
|
||||||
|
/** Adds {@code updatedAt} so the eviction sweep can find stale terminals. */
|
||||||
|
public record Session(String sessionId, Long keyId, Long kbId, String topic, Status status,
|
||||||
|
ResearchResult result, String error, Instant updatedAt) {
|
||||||
|
|
||||||
|
/** Convenience for {@link #register} (status=RUNNING, no result). */
|
||||||
|
static Session running(String sessionId, Long keyId, Long kbId, String topic) {
|
||||||
|
return new Session(sessionId, keyId, kbId, topic, Status.RUNNING, null, null, Instant.now());
|
||||||
|
}
|
||||||
|
|
||||||
|
private Session with(Status newStatus, ResearchResult res, String err) {
|
||||||
|
return new Session(sessionId, keyId, kbId, topic, newStatus, res, err, Instant.now());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Exception thrown by {@link #startIfAllowed} when the per-key cap is hit. */
|
||||||
|
public static class TooManyConcurrentException extends RuntimeException {
|
||||||
|
public TooManyConcurrentException(String msg) { super(msg); }
|
||||||
|
}
|
||||||
|
|
||||||
|
private final Map<String, Session> sessions = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Per-key count of RUNNING sessions, kept in lock-step with the
|
||||||
|
* {@code status==RUNNING} sessions in {@link #sessions}. Maintained
|
||||||
|
* atomically so {@link #startIfAllowed} can enforce the cap without a
|
||||||
|
* check-then-act race (two concurrent starts could both pass a stream-based
|
||||||
|
* count and both put). Incremented on start, decremented on each
|
||||||
|
* RUNNING→terminal transition (complete/fail/cancel).
|
||||||
|
*/
|
||||||
|
private final Map<Long, AtomicInteger> runningPerKey = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
|
private final int maxConcurrentPerKey;
|
||||||
|
private final Duration sessionTtl;
|
||||||
|
|
||||||
|
public KbResearchSessionRegistry(
|
||||||
|
@Value("${mate.kbopen.research.max-concurrent-per-key:3}") int maxConcurrentPerKey,
|
||||||
|
@Value("${mate.kbopen.research.session-ttl:PT30M}") Duration sessionTtl) {
|
||||||
|
this.maxConcurrentPerKey = maxConcurrentPerKey;
|
||||||
|
this.sessionTtl = sessionTtl;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Reserve a slot for a new session, enforcing the per-key concurrency cap.
|
||||||
|
*
|
||||||
|
* <p>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<Session> get(String sessionId) {
|
||||||
|
return Optional.ofNullable(sessions.get(sessionId));
|
||||||
|
}
|
||||||
|
|
||||||
|
/** RUNNING → COMPLETED. No-op on a session that was already CANCELLED (sticky terminal). */
|
||||||
|
public void complete(String sessionId, ResearchResult result) {
|
||||||
|
sessions.computeIfPresent(sessionId, (k, s) -> {
|
||||||
|
if (s.status() == Status.CANCELLED) {
|
||||||
|
return s; // sticky terminal — no transition, no counter change
|
||||||
|
}
|
||||||
|
decrementRunning(s.keyId()); // RUNNING → COMPLETED releases the slot
|
||||||
|
return s.with(Status.COMPLETED, result, null);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/** RUNNING → FAILED. No-op on a session that was already CANCELLED (sticky terminal). */
|
||||||
|
public void fail(String sessionId, String error) {
|
||||||
|
sessions.computeIfPresent(sessionId, (k, s) -> {
|
||||||
|
if (s.status() == Status.CANCELLED) {
|
||||||
|
return s;
|
||||||
|
}
|
||||||
|
decrementRunning(s.keyId());
|
||||||
|
return s.with(Status.FAILED, null, error);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/** RUNNING → CANCELLED. Returns false if the session is missing or already terminal. */
|
||||||
|
public boolean cancel(String sessionId) {
|
||||||
|
Session[] before = new Session[1];
|
||||||
|
sessions.computeIfPresent(sessionId, (k, s) -> {
|
||||||
|
before[0] = s;
|
||||||
|
if (s.status() == Status.RUNNING) {
|
||||||
|
decrementRunning(s.keyId());
|
||||||
|
return s.with(Status.CANCELLED, null, null);
|
||||||
|
}
|
||||||
|
return s;
|
||||||
|
});
|
||||||
|
return before[0] != null && before[0].status() == Status.RUNNING;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Release one running-slot for {@code keyId}, floored at 0. */
|
||||||
|
private void decrementRunning(Long keyId) {
|
||||||
|
AtomicInteger count = runningPerKey.get(keyId);
|
||||||
|
if (count != null) {
|
||||||
|
// getAndDeccrement would go negative; clamp instead so repeated
|
||||||
|
// terminal transitions (e.g. complete after cancel) can't drift.
|
||||||
|
while (true) {
|
||||||
|
int cur = count.get();
|
||||||
|
if (cur <= 0) break;
|
||||||
|
if (count.compareAndSet(cur, cur - 1)) break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Drop terminal sessions older than {@code sessionTtl}. Called periodically
|
||||||
|
* by {@link #evictExpired}; public for testing.
|
||||||
|
*/
|
||||||
|
public int evictExpired(Instant now) {
|
||||||
|
int removed = 0;
|
||||||
|
for (Map.Entry<String, Session> e : sessions.entrySet()) {
|
||||||
|
Session s = e.getValue();
|
||||||
|
if (s.status() != Status.RUNNING && now.isAfter(s.updatedAt().plus(sessionTtl))) {
|
||||||
|
if (sessions.remove(e.getKey()) != null) removed++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (removed > 0) {
|
||||||
|
log.info("[KbResearchSessionRegistry] Evicted {} terminal session(s) older than {}", removed, sessionTtl);
|
||||||
|
}
|
||||||
|
return removed;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Scheduled sweep — runs every 5 min. */
|
||||||
|
@Scheduled(fixedDelay = 5 * 60 * 1000L)
|
||||||
|
public void evictExpired() {
|
||||||
|
evictExpired(Instant.now());
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -81,6 +81,10 @@ public class WikiResearchService {
|
|||||||
broadcast(sessionId, "research.plan", Map.of(
|
broadcast(sessionId, "research.plan", Map.of(
|
||||||
"questions", questions.stream().map(q -> Map.of("question", q.question, "intent", q.intent)).toList()
|
"questions", questions.stream().map(q -> Map.of("question", q.question, "intent", q.intent)).toList()
|
||||||
));
|
));
|
||||||
|
// Cooperative cancellation: the Open API cancel endpoint calls
|
||||||
|
// streamTracker.requestStop(sessionId). Bail before the expensive
|
||||||
|
// retrieve+draft fan-out so cancel actually halts LLM cost.
|
||||||
|
ensureNotCancelled(sessionId);
|
||||||
|
|
||||||
// Stage 2: Retrieve + Draft (并行)
|
// Stage 2: Retrieve + Draft (并行)
|
||||||
List<Section> sections = draftStage(kbId, questions, topK, sessionId);
|
List<Section> sections = draftStage(kbId, questions, topK, sessionId);
|
||||||
@ -88,6 +92,8 @@ public class WikiResearchService {
|
|||||||
broadcast(sessionId, "research.error", Map.of("message", i18n.msg("research.broadcast.draft_all_empty")));
|
broadcast(sessionId, "research.error", Map.of("message", i18n.msg("research.broadcast.draft_all_empty")));
|
||||||
return new ResearchResult(topic, sections, i18n.msg("research.fallback.no_materials"));
|
return new ResearchResult(topic, sections, i18n.msg("research.fallback.no_materials"));
|
||||||
}
|
}
|
||||||
|
// Second checkpoint before the compose LLM call.
|
||||||
|
ensureNotCancelled(sessionId);
|
||||||
|
|
||||||
// Stage 3: Compose
|
// Stage 3: Compose
|
||||||
String report = composeStage(topic, sections);
|
String report = composeStage(topic, sections);
|
||||||
@ -97,6 +103,11 @@ public class WikiResearchService {
|
|||||||
"materialsUsed", sections.stream().flatMap(s -> s.materialRefs.stream()).distinct().count()
|
"materialsUsed", sections.stream().flatMap(s -> s.materialRefs.stream()).distinct().count()
|
||||||
));
|
));
|
||||||
return new ResearchResult(topic, sections, report);
|
return new ResearchResult(topic, sections, report);
|
||||||
|
} catch (ResearchCancelledException ce) {
|
||||||
|
// Expected: caller has already flipped the session to CANCELLED
|
||||||
|
// and closed the SSE stream. Do not broadcast an error event.
|
||||||
|
log.info("[Research] Cancelled: kbId={}, topic={}, sessionId={}", kbId, topic, sessionId);
|
||||||
|
return new ResearchResult(topic, List.of(), "Research cancelled by user");
|
||||||
} catch (Exception e) {
|
} catch (Exception e) {
|
||||||
log.error("[Research] Failed: kbId={}, topic={}: {}", kbId, topic, e.getMessage(), e);
|
log.error("[Research] Failed: kbId={}, topic={}: {}", kbId, topic, e.getMessage(), e);
|
||||||
broadcast(sessionId, "research.error", Map.of("message", e.getMessage() != null ? e.getMessage() : i18n.msg("research.broadcast.failed")));
|
broadcast(sessionId, "research.error", Map.of("message", e.getMessage() != null ? e.getMessage() : i18n.msg("research.broadcast.failed")));
|
||||||
@ -104,6 +115,17 @@ public class WikiResearchService {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Throws {@link ResearchCancelledException} if the caller (via the Open API
|
||||||
|
* cancel endpoint) has called {@link ChatStreamTracker#requestStop} on this
|
||||||
|
* session. Checked at each pipeline stage boundary.
|
||||||
|
*/
|
||||||
|
private void ensureNotCancelled(String sessionId) {
|
||||||
|
if (streamTracker.isStopRequested(sessionId)) {
|
||||||
|
throw new ResearchCancelledException(sessionId);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ==================== Stage 1: Plan ====================
|
// ==================== Stage 1: Plan ====================
|
||||||
|
|
||||||
private List<SubQuestion> planStage(String topic) {
|
private List<SubQuestion> planStage(String topic) {
|
||||||
@ -151,6 +173,12 @@ public class WikiResearchService {
|
|||||||
Thread.currentThread().interrupt();
|
Thread.currentThread().interrupt();
|
||||||
return new Section(q.question, "", List.of());
|
return new Section(q.question, "", List.of());
|
||||||
}
|
}
|
||||||
|
// Re-check cancellation inside the parallel draft loop too —
|
||||||
|
// draftOneSection issues its own LLM call, which is the main
|
||||||
|
// cost driver, so skip queued-but-not-started drafts on cancel.
|
||||||
|
if (streamTracker.isStopRequested(sessionId)) {
|
||||||
|
return new Section(q.question, "", List.of());
|
||||||
|
}
|
||||||
try {
|
try {
|
||||||
Section section = draftOneSection(kbId, q, topK);
|
Section section = draftOneSection(kbId, q, topK);
|
||||||
broadcast(sessionId, "research.draft", Map.of(
|
broadcast(sessionId, "research.draft", Map.of(
|
||||||
@ -300,4 +328,16 @@ public class WikiResearchService {
|
|||||||
public record Section(String question, String content, List<MaterialRef> materialRefs) {}
|
public record Section(String question, String content, List<MaterialRef> materialRefs) {}
|
||||||
|
|
||||||
public record ResearchResult(String topic, List<Section> sections, String report) {}
|
public record ResearchResult(String topic, List<Section> sections, String report) {}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Thrown when {@link ChatStreamTracker#requestStop(String)} was called on
|
||||||
|
* the session between pipeline stages. The {@link #research} method catches
|
||||||
|
* this to short-circuit — the caller (the Open API controller) has already
|
||||||
|
* flipped the registry to CANCELLED and closed the SSE stream.
|
||||||
|
*/
|
||||||
|
public static class ResearchCancelledException extends RuntimeException {
|
||||||
|
public ResearchCancelledException(String sessionId) {
|
||||||
|
super("Research cancelled: sessionId=" + sessionId);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@ -0,0 +1,126 @@
|
|||||||
|
package vip.mate.kbopen.auth;
|
||||||
|
|
||||||
|
import org.junit.jupiter.api.BeforeEach;
|
||||||
|
import org.junit.jupiter.api.DisplayName;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.springframework.mock.web.MockFilterChain;
|
||||||
|
import org.springframework.mock.web.MockHttpServletRequest;
|
||||||
|
import org.springframework.mock.web.MockHttpServletResponse;
|
||||||
|
import vip.mate.kbopen.auth.KbApiKeyService.AuthResult;
|
||||||
|
|
||||||
|
import java.time.LocalDateTime;
|
||||||
|
import java.util.Optional;
|
||||||
|
import java.util.Set;
|
||||||
|
|
||||||
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
|
import static org.mockito.ArgumentMatchers.anyInt;
|
||||||
|
import static org.mockito.ArgumentMatchers.anyLong;
|
||||||
|
import static org.mockito.Mockito.mock;
|
||||||
|
import static org.mockito.Mockito.never;
|
||||||
|
import static org.mockito.Mockito.verify;
|
||||||
|
import static org.mockito.Mockito.when;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Unit tests for {@link KbOpenApiAuthFilter} — focuses on the R7 SSE token
|
||||||
|
* fallback and R5 scope limitation added in #446:
|
||||||
|
* <ul>
|
||||||
|
* <li>{@code ?token=} is accepted on SSE stream paths (EventSource can't set
|
||||||
|
* an Authorization header).</li>
|
||||||
|
* <li>{@code ?token=} is rejected on non-SSE paths so the key doesn't leak
|
||||||
|
* into access / proxy logs.</li>
|
||||||
|
* <li>The SSE stream path bypasses the per-minute rate limiter (reconnects /
|
||||||
|
* heartbeats must not burn the key's start quota).</li>
|
||||||
|
* </ul>
|
||||||
|
*/
|
||||||
|
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());
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -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<Session> session = registry.get("s1");
|
||||||
|
assertThat(session).isPresent();
|
||||||
|
assertThat(session.get().status()).isEqualTo(Status.RUNNING);
|
||||||
|
assertThat(session.get().keyId()).isEqualTo(100L);
|
||||||
|
assertThat(session.get().kbId()).isEqualTo(10L);
|
||||||
|
assertThat(session.get().topic()).isEqualTo("test topic");
|
||||||
|
assertThat(session.get().updatedAt()).isNotNull();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("complete transitions RUNNING → COMPLETED with result")
|
||||||
|
void completeTransitionsToCompleted() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
|
||||||
|
registry.complete("s1", RESULT);
|
||||||
|
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.COMPLETED);
|
||||||
|
assertThat(session.result()).isEqualTo(RESULT);
|
||||||
|
assertThat(session.result().report()).isEqualTo("final report");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("fail transitions RUNNING → FAILED with error")
|
||||||
|
void failTransitionsToFailed() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
|
||||||
|
registry.fail("s1", "LLM timeout");
|
||||||
|
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.FAILED);
|
||||||
|
assertThat(session.error()).isEqualTo("LLM timeout");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("cancel transitions RUNNING → CANCELLED")
|
||||||
|
void cancelTransitionsToCancelled() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
|
||||||
|
boolean cancelled = registry.cancel("s1");
|
||||||
|
|
||||||
|
assertThat(cancelled).isTrue();
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.CANCELLED);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("cancel on non-running session returns false (no-op)")
|
||||||
|
void cancelOnCompletedIsNoop() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
registry.complete("s1", RESULT);
|
||||||
|
|
||||||
|
boolean cancelled = registry.cancel("s1");
|
||||||
|
|
||||||
|
assertThat(cancelled).isFalse();
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.COMPLETED);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("get on unknown session returns empty")
|
||||||
|
void getUnknownReturnsEmpty() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
assertThat(registry.get("nonexistent")).isEmpty();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Review #446: sticky CANCELLED terminal ────────────────────────────
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("complete after cancel is a no-op — CANCELLED is sticky")
|
||||||
|
void completeAfterCancelIsNoop() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
registry.cancel("s1");
|
||||||
|
|
||||||
|
// Late complete() arriving from the async pipeline must NOT overwrite
|
||||||
|
// the CANCELLED terminal the user explicitly requested.
|
||||||
|
registry.complete("s1", RESULT);
|
||||||
|
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.CANCELLED);
|
||||||
|
assertThat(session.result()).isNull();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("fail after cancel is a no-op — CANCELLED is sticky")
|
||||||
|
void failAfterCancelIsNoop() {
|
||||||
|
KbResearchSessionRegistry registry = newRegistry();
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "topic");
|
||||||
|
registry.cancel("s1");
|
||||||
|
|
||||||
|
registry.fail("s1", "race condition");
|
||||||
|
|
||||||
|
Session session = registry.get("s1").get();
|
||||||
|
assertThat(session.status()).isEqualTo(Status.CANCELLED);
|
||||||
|
assertThat(session.error()).isNull();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Review #446: per-key concurrency cap ──────────────────────────────
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("startIfAllowed throws once the per-key cap is reached")
|
||||||
|
void startIfAllowedEnforcesCap() {
|
||||||
|
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(2, Duration.ofMinutes(30));
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "t1");
|
||||||
|
registry.startIfAllowed("s2", 100L, 10L, "t2");
|
||||||
|
|
||||||
|
// Third running session for the same key should be rejected → 429 upstream.
|
||||||
|
assertThatThrownBy(() -> registry.startIfAllowed("s3", 100L, 10L, "t3"))
|
||||||
|
.isInstanceOf(TooManyConcurrentException.class)
|
||||||
|
.hasMessageContaining("limit is 2");
|
||||||
|
|
||||||
|
// A different key is unaffected (cap is per-key, not global).
|
||||||
|
registry.startIfAllowed("s4", 200L, 10L, "t4");
|
||||||
|
assertThat(registry.get("s4")).isPresent();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("completed sessions do not count toward the running cap")
|
||||||
|
void completedDoesNotCountTowardCap() {
|
||||||
|
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(1, Duration.ofMinutes(30));
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "t1");
|
||||||
|
registry.complete("s1", RESULT);
|
||||||
|
|
||||||
|
// The terminal session no longer occupies a slot.
|
||||||
|
registry.startIfAllowed("s2", 100L, 10L, "t2");
|
||||||
|
assertThat(registry.get("s2")).isPresent();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("cancelled and failed sessions also release their slot (counter consistency)")
|
||||||
|
void cancelledAndFailedReleaseSlot() {
|
||||||
|
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(1, Duration.ofMinutes(30));
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "t1");
|
||||||
|
registry.cancel("s1"); // RUNNING → CANCELLED releases slot
|
||||||
|
registry.startIfAllowed("s2", 100L, 10L, "t2");
|
||||||
|
assertThat(registry.get("s2")).isPresent();
|
||||||
|
|
||||||
|
registry.fail("s2", "boom"); // RUNNING → FAILED releases slot
|
||||||
|
registry.startIfAllowed("s3", 100L, 10L, "t3");
|
||||||
|
assertThat(registry.get("s3")).isPresent();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("concurrent starts never exceed the per-key cap (no check-then-act race)")
|
||||||
|
void startIfAllowedIsAtomicUnderConcurrency() throws Exception {
|
||||||
|
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<>();
|
||||||
|
|
||||||
|
for (int i = 0; i < threads; i++) {
|
||||||
|
String sid = "concurrent-" + i;
|
||||||
|
Thread t = Thread.ofVirtual().start(() -> {
|
||||||
|
try {
|
||||||
|
start.await();
|
||||||
|
registry.startIfAllowed(sid, 100L, 10L, "t");
|
||||||
|
admitted.incrementAndGet();
|
||||||
|
} catch (TooManyConcurrentException e) {
|
||||||
|
rejected.incrementAndGet();
|
||||||
|
} catch (InterruptedException e) {
|
||||||
|
Thread.currentThread().interrupt();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
workers.add(t);
|
||||||
|
}
|
||||||
|
start.countDown();
|
||||||
|
for (Thread t : workers) t.join();
|
||||||
|
|
||||||
|
// The whole point: exactly `cap` sessions get in, no matter the
|
||||||
|
// scheduling. The old stream-and-count impl could admit more under
|
||||||
|
// contention.
|
||||||
|
assertThat(admitted.get()).isEqualTo(cap);
|
||||||
|
assertThat(rejected.get()).isEqualTo(threads - cap);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Review #446: TTL eviction ─────────────────────────────────────────
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("evictExpired removes terminal sessions past TTL but keeps RUNNING + fresh terminals")
|
||||||
|
void evictExpiredRemovesStaleTerminals() {
|
||||||
|
// Tiny TTL so a fresh terminal (updatedAt≈now) is clearly within window.
|
||||||
|
Duration ttl = Duration.ofSeconds(60);
|
||||||
|
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(3, ttl);
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "running"); // RUNNING — never evicted
|
||||||
|
registry.startIfAllowed("s2", 100L, 10L, "stale-completed");
|
||||||
|
registry.complete("s2", RESULT); // terminal, will be aged
|
||||||
|
registry.startIfAllowed("s3", 100L, 10L, "fresh-cancelled");
|
||||||
|
registry.cancel("s3"); // terminal, fresh
|
||||||
|
|
||||||
|
// Horizon well past TTL: both terminals s2 and s3 are now stale.
|
||||||
|
Instant far = Instant.now().plus(Duration.ofSeconds(120));
|
||||||
|
int removed = registry.evictExpired(far);
|
||||||
|
assertThat(removed).isEqualTo(2);
|
||||||
|
assertThat(registry.get("s1")).isPresent(); // running always kept
|
||||||
|
assertThat(registry.get("s2")).isEmpty();
|
||||||
|
assertThat(registry.get("s3")).isEmpty();
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
@DisplayName("evictExpired keeps a fresh terminal within the TTL window")
|
||||||
|
void evictExpiredKeepsFreshTerminal() {
|
||||||
|
Duration ttl = Duration.ofHours(1);
|
||||||
|
KbResearchSessionRegistry registry = new KbResearchSessionRegistry(3, ttl);
|
||||||
|
registry.startIfAllowed("s1", 100L, 10L, "just-completed");
|
||||||
|
registry.complete("s1", RESULT); // updatedAt ≈ now
|
||||||
|
|
||||||
|
// Evict only 5 seconds later — well within the 1h TTL.
|
||||||
|
int removed = registry.evictExpired(Instant.now().plus(Duration.ofSeconds(5)));
|
||||||
|
assertThat(removed).isZero();
|
||||||
|
assertThat(registry.get("s1")).isPresent();
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -252,7 +252,7 @@ KbOpenApiAuthFilter (OncePerRequestFilter, 仅拦截 /api/v1/open/kb/**)
|
|||||||
|
|
||||||
| Scope | 允许的操作 | 说明 |
|
| Scope | 允许的操作 | 说明 |
|
||||||
|---|---|---|
|
|---|---|---|
|
||||||
| `kb:search` | POST `/search`、POST `/search/chunks` | 检索 |
|
| `kb:search` | POST `/search`、POST `/search/chunks`、POST `/research/**`（Deep Research 全套）| 检索（含异步深度研究，复用同一 scope）|
|
||||||
| `kb:read` | GET `/pages/{slug}`、GET `/pages/{slug}/trace`、POST `/pages/{slug}/traverse` | 读主体画像 + 溯源 + 关系遍历 |
|
| `kb:read` | GET `/pages/{slug}`、GET `/pages/{slug}/trace`、POST `/pages/{slug}/traverse` | 读主体画像 + 溯源 + 关系遍历 |
|
||||||
| `kb:list` | GET `/pages`、GET `/taxonomy` | 列页面 + 地图 |
|
| `kb:list` | GET `/pages`、GET `/taxonomy` | 列页面 + 地图 |
|
||||||
| `kb:meta` | GET `/stats`、GET `/whats-new` | 元信息 + 时效查询 |
|
| `kb:meta` | GET `/stats`、GET `/whats-new` | 元信息 + 时效查询 |
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user