001/* 002 * Copyright 2022-2026 Revetware LLC. 003 * 004 * Licensed under the Apache License, Version 2.0 (the "License"); 005 * you may not use this file except in compliance with the License. 006 * You may obtain a copy of the License at 007 * 008 * http://www.apache.org/licenses/LICENSE-2.0 009 * 010 * Unless required by applicable law or agreed to in writing, software 011 * distributed under the License is distributed on an "AS IS" BASIS, 012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 013 * See the License for the specific language governing permissions and 014 * limitations under the License. 015 */ 016 017package com.soklet; 018 019import org.jspecify.annotations.NonNull; 020import org.jspecify.annotations.Nullable; 021 022import javax.annotation.concurrent.NotThreadSafe; 023import javax.annotation.concurrent.ThreadSafe; 024import java.time.Duration; 025import java.time.Instant; 026import java.util.ArrayList; 027import java.util.List; 028import java.util.Optional; 029import java.util.concurrent.ConcurrentHashMap; 030import java.util.concurrent.ConcurrentMap; 031import java.util.concurrent.atomic.AtomicInteger; 032import java.util.function.Predicate; 033 034import static java.time.Duration.ZERO; 035import static java.time.Duration.between; 036import static java.time.Duration.ofHours; 037import static java.time.Duration.ofMinutes; 038import static java.util.Objects.requireNonNull; 039 040/** 041 * Compare-and-set persistence contract for MCP sessions. 042 * 043 * @author <a href="https://www.revetkn.com">Mark Allen</a> 044 */ 045@ThreadSafe 046public interface McpSessionStore { 047 /** 048 * Creates and admits a new session for the given request and endpoint class. 049 * <p> 050 * Implementations are responsible for generating a valid MCP session ID and 051 * enforcing their own concurrency limits atomically with persistence. Returning 052 * {@link Optional#empty()} declines admission before session state is created 053 * (for example, when a concurrent-session limit is reached). 054 * <p> 055 * Returned sessions must be new, uninitialized, unterminated sessions for 056 * {@code endpointClass}. Soklet will complete initialization with 057 * {@link #replace(McpStoredSession, McpStoredSession)} after the endpoint's 058 * {@link McpEndpoint#initialize(McpInitializationContext, McpSessionContext)} 059 * callback succeeds. 060 * 061 * @param request the initialize request 062 * @param endpointClass the MCP endpoint class that will own the session 063 * @return the newly-created stored session, or empty if the store declined admission 064 */ 065 @NonNull 066 Optional<McpStoredSession> create(@NonNull Request request, 067 @NonNull Class<? extends McpEndpoint> endpointClass); 068 069 /** 070 * Loads a session by ID. 071 * 072 * @param sessionId the session ID 073 * @return the stored session, if it exists and has not expired 074 */ 075 @NonNull 076 Optional<McpStoredSession> findBySessionId(@NonNull String sessionId); 077 078 /** 079 * Replaces a session using compare-and-set semantics. 080 * 081 * @param expected the currently-stored session snapshot 082 * @param updated the replacement session snapshot 083 * @return {@code true} if the replacement succeeded 084 */ 085 @NonNull 086 Boolean replace(@NonNull McpStoredSession expected, 087 @NonNull McpStoredSession updated); 088 089 /** 090 * Deletes a session by ID. 091 * 092 * @param sessionId the session ID to delete 093 */ 094 void deleteBySessionId(@NonNull String sessionId); 095 096 /** 097 * Acquires a builder for Soklet's default in-memory MCP session store. 098 * 099 * @return a new MCP session store builder 100 */ 101 @NonNull 102 static Builder builder() { 103 return new Builder(); 104 } 105 106 /** 107 * Acquires Soklet's default in-memory MCP session store. 108 * 109 * @return a new in-memory session store 110 */ 111 @NonNull 112 static McpSessionStore fromDefaults() { 113 return builder().build(); 114 } 115 116 /** 117 * Acquires the default in-memory session store using Soklet's default idle timeout. 118 * <p> 119 * Expired sessions are reclaimed opportunistically during lookup and subsequent 120 * session-creation activity; exact deletion timing is therefore best-effort 121 * rather than timer-driven. 122 * 123 * @return a new in-memory session store 124 */ 125 @NonNull 126 static McpSessionStore fromInMemory() { 127 return fromDefaults(); 128 } 129 130 /** 131 * Builder for Soklet's default in-memory MCP session store. 132 */ 133 @NotThreadSafe 134 final class Builder { 135 @Nullable 136 Duration idleTimeout; 137 @Nullable 138 IdGenerator<String> sessionIdGenerator; 139 @Nullable 140 Integer concurrentSessionLimit; 141 142 private Builder() {} 143 144 /** 145 * Sets the idle timeout, or {@link Duration#ZERO} to disable idle expiry. 146 * <p> 147 * <strong>Warning:</strong> with idle expiry disabled, sessions only release their slots via 148 * explicit termination or {@link McpSessionStore#deleteBySessionId(String)}. If you also 149 * configure a finite {@link #concurrentSessionLimit(Integer)}, abandoned sessions (e.g. clients 150 * that disappear without terminating) will occupy slots indefinitely and can eventually block 151 * all new session admission. Ensure your application enforces session lifecycle cleanup, or 152 * keep a nonzero idle timeout. 153 * 154 * @param idleTimeout the idle timeout, or {@code null} for the default 155 * @return this builder 156 */ 157 @NonNull 158 public Builder idleTimeout(@Nullable Duration idleTimeout) { 159 this.idleTimeout = idleTimeout; 160 return this; 161 } 162 163 /** 164 * Sets the generator used for newly-created MCP session IDs. 165 * <p> 166 * Custom generators must return globally unique, cryptographically strong, 167 * visible-ASCII IDs suitable for {@code MCP-Session-Id} header values. 168 * 169 * @param sessionIdGenerator the session ID generator, or {@code null} for the default 170 * @return this builder 171 */ 172 @NonNull 173 public Builder sessionIdGenerator(@Nullable IdGenerator<String> sessionIdGenerator) { 174 this.sessionIdGenerator = sessionIdGenerator; 175 return this; 176 } 177 178 /** 179 * Sets the concurrent MCP session limit. 180 * <p> 181 * A value of {@code 0} disables the in-memory store's session cap. 182 * 183 * @param concurrentSessionLimit the concurrent MCP session limit, or {@code null} for the default 184 * @return this builder 185 */ 186 @NonNull 187 public Builder concurrentSessionLimit(@Nullable Integer concurrentSessionLimit) { 188 this.concurrentSessionLimit = concurrentSessionLimit; 189 return this; 190 } 191 192 /** 193 * Builds the MCP session store. 194 * 195 * @return the built MCP session store 196 */ 197 @NonNull 198 public McpSessionStore build() { 199 return new DefaultMcpSessionStore(this); 200 } 201 } 202} 203 204final class DefaultMcpSessionStore implements McpSessionStore { 205 @NonNull 206 private static final Duration DEFAULT_IDLE_TIMEOUT; 207 @NonNull 208 private static final Duration DEFAULT_SWEEP_INTERVAL; 209 @NonNull 210 private static final Integer DEFAULT_CONCURRENT_SESSION_LIMIT; 211 @NonNull 212 private final Duration idleTimeout; 213 @NonNull 214 private final IdGenerator<String> sessionIdGenerator; 215 @NonNull 216 private final Integer concurrentSessionLimit; 217 @NonNull 218 private final ConcurrentMap<String, McpStoredSession> sessions; 219 @NonNull 220 private final ConcurrentMap<String, Boolean> activeLimitedSessionIds; 221 @NonNull 222 private final AtomicInteger activeLimitedSessionCount; 223 @NonNull 224 private volatile Predicate<String> pinnedSessionPredicate; 225 @NonNull 226 private volatile Instant lastSweepAt; 227 228 static { 229 DEFAULT_IDLE_TIMEOUT = ofHours(24); 230 DEFAULT_SWEEP_INTERVAL = ofMinutes(1); 231 DEFAULT_CONCURRENT_SESSION_LIMIT = 8_192; 232 } 233 234 DefaultMcpSessionStore(McpSessionStore.@NonNull Builder builder) { 235 requireNonNull(builder); 236 this.idleTimeout = builder.idleTimeout != null ? builder.idleTimeout : DEFAULT_IDLE_TIMEOUT; 237 this.sessionIdGenerator = builder.sessionIdGenerator != null ? builder.sessionIdGenerator : IdGenerator.defaultSessionInstance(); 238 this.concurrentSessionLimit = builder.concurrentSessionLimit != null ? builder.concurrentSessionLimit : DEFAULT_CONCURRENT_SESSION_LIMIT; 239 this.sessions = new ConcurrentHashMap<>(); 240 this.activeLimitedSessionIds = new ConcurrentHashMap<>(); 241 this.activeLimitedSessionCount = new AtomicInteger(0); 242 this.pinnedSessionPredicate = sessionId -> false; 243 this.lastSweepAt = Instant.EPOCH; 244 245 if (this.idleTimeout.isNegative()) 246 throw new IllegalArgumentException("Idle timeout must not be negative."); 247 248 if (this.concurrentSessionLimit < 0) 249 throw new IllegalArgumentException("Concurrent session limit must be >= 0"); 250 } 251 252 @NonNull 253 @Override 254 public synchronized Optional<McpStoredSession> create(@NonNull Request request, 255 @NonNull Class<? extends McpEndpoint> endpointClass) { 256 requireNonNull(request); 257 requireNonNull(endpointClass); 258 takeExpiredSessionsIfSweepDue(); 259 260 if (!reserveSessionSlot()) 261 return Optional.empty(); 262 263 String sessionId = this.sessionIdGenerator.generateId(request); 264 265 if (!DefaultMcpRuntime.isValidMcpSessionId(sessionId)) { 266 releaseReservedSessionSlot(); 267 throw new IllegalStateException("MCP session ID generator produced an invalid session ID."); 268 } 269 270 takeExpiredSession(sessionId); 271 272 Instant now = Instant.now(); 273 McpStoredSession session = new McpStoredSession( 274 sessionId, 275 endpointClass, 276 now, 277 now, 278 false, 279 false, 280 null, 281 null, 282 null, 283 McpSessionContext.fromBlankSlate(), 284 null, 285 0L 286 ); 287 288 try { 289 put(session); 290 } catch (Throwable throwable) { 291 releaseReservedSessionSlot(); 292 throw throwable; 293 } 294 295 this.activeLimitedSessionIds.put(sessionId, Boolean.TRUE); 296 return Optional.of(session); 297 } 298 299 synchronized void create(@NonNull McpStoredSession session) { 300 requireNonNull(session); 301 takeExpiredSessionsIfSweepDue(); 302 303 boolean slotReserved = false; 304 305 if (session.terminatedAt() == null) { 306 if (!reserveSessionSlot()) 307 throw new IllegalStateException("MCP session limit reached."); 308 309 slotReserved = true; 310 } 311 312 try { 313 put(session); 314 } catch (Throwable throwable) { 315 if (slotReserved) 316 releaseReservedSessionSlot(); 317 318 throw throwable; 319 } 320 321 if (slotReserved) 322 this.activeLimitedSessionIds.put(session.sessionId(), Boolean.TRUE); 323 } 324 325 @NonNull 326 @Override 327 public synchronized Optional<McpStoredSession> findBySessionId(@NonNull String sessionId) { 328 requireNonNull(sessionId); 329 330 McpStoredSession storedSession = this.sessions.get(sessionId); 331 332 if (storedSession == null) 333 return Optional.empty(); 334 335 if (isExpired(storedSession)) { 336 if (this.sessions.remove(sessionId, storedSession)) 337 releaseSessionSlot(sessionId); 338 339 return Optional.empty(); 340 } 341 342 return Optional.of(storedSession); 343 } 344 345 @NonNull 346 @Override 347 public synchronized Boolean replace(@NonNull McpStoredSession expected, 348 @NonNull McpStoredSession updated) { 349 requireNonNull(expected); 350 requireNonNull(updated); 351 352 if (!expected.sessionId().equals(updated.sessionId())) 353 throw new IllegalArgumentException("Expected and updated sessions must have the same session ID."); 354 355 if (updated.version().longValue() <= expected.version().longValue()) 356 throw new IllegalArgumentException("Updated session version must be strictly greater than expected version."); 357 358 if (isExpired(expected)) 359 return false; 360 361 boolean replaced = this.sessions.replace(expected.sessionId(), expected, updated); 362 363 if (replaced && expected.terminatedAt() == null && updated.terminatedAt() != null) 364 releaseSessionSlot(updated.sessionId()); 365 366 return replaced; 367 } 368 369 @Override 370 public synchronized void deleteBySessionId(@NonNull String sessionId) { 371 requireNonNull(sessionId); 372 373 if (this.sessions.remove(sessionId) != null) 374 releaseSessionSlot(sessionId); 375 } 376 377 synchronized void pinnedSessionPredicate(@NonNull Predicate<String> pinnedSessionPredicate) { 378 requireNonNull(pinnedSessionPredicate); 379 this.pinnedSessionPredicate = pinnedSessionPredicate; 380 } 381 382 synchronized boolean containsSessionId(@NonNull String sessionId) { 383 requireNonNull(sessionId); 384 return this.sessions.containsKey(sessionId); 385 } 386 387 @NonNull 388 synchronized Optional<McpStoredSession> takeExpiredSession(@NonNull String sessionId) { 389 requireNonNull(sessionId); 390 391 McpStoredSession storedSession = this.sessions.get(sessionId); 392 393 if (storedSession == null || !isExpired(storedSession)) 394 return Optional.empty(); 395 396 if (!this.sessions.remove(sessionId, storedSession)) 397 return Optional.empty(); 398 399 releaseSessionSlot(sessionId); 400 return Optional.of(storedSession); 401 } 402 403 @NonNull 404 synchronized List<McpStoredSession> takeExpiredSessionsIfSweepDue() { 405 if (ZERO.equals(this.idleTimeout)) 406 return List.of(); 407 408 Instant now = Instant.now(); 409 Duration sweepInterval = sweepInterval(); 410 411 if (between(this.lastSweepAt, now).compareTo(sweepInterval) < 0) 412 return List.of(); 413 414 this.lastSweepAt = now; 415 416 List<McpStoredSession> expiredSessions = new ArrayList<>(); 417 418 for (var entry : this.sessions.entrySet()) { 419 McpStoredSession storedSession = entry.getValue(); 420 421 if (isExpired(storedSession) && this.sessions.remove(entry.getKey(), storedSession)) { 422 expiredSessions.add(storedSession); 423 releaseSessionSlot(entry.getKey()); 424 } 425 } 426 427 return expiredSessions; 428 } 429 430 private void put(@NonNull McpStoredSession session) { 431 requireNonNull(session); 432 433 McpStoredSession previous = this.sessions.putIfAbsent(session.sessionId(), session); 434 435 if (previous != null) 436 throw new IllegalStateException("Session with ID '%s' already exists".formatted(session.sessionId())); 437 } 438 439 private boolean reserveSessionSlot() { 440 if (this.concurrentSessionLimit == 0) 441 return true; 442 443 while (true) { 444 int current = this.activeLimitedSessionCount.get(); 445 446 if (current >= this.concurrentSessionLimit) 447 return false; 448 449 if (this.activeLimitedSessionCount.compareAndSet(current, current + 1)) 450 return true; 451 } 452 } 453 454 private void releaseSessionSlot(@NonNull String sessionId) { 455 requireNonNull(sessionId); 456 457 if (this.activeLimitedSessionIds.remove(sessionId) != null) 458 releaseReservedSessionSlot(); 459 } 460 461 private void releaseReservedSessionSlot() { 462 while (true) { 463 int current = this.activeLimitedSessionCount.get(); 464 465 if (current <= 0) 466 return; 467 468 if (this.activeLimitedSessionCount.compareAndSet(current, current - 1)) 469 return; 470 } 471 } 472 473 private boolean isExpired(@NonNull McpStoredSession storedSession) { 474 requireNonNull(storedSession); 475 476 if (ZERO.equals(this.idleTimeout)) 477 return false; 478 479 if (storedSession.terminatedAt() != null) 480 return false; 481 482 if (this.pinnedSessionPredicate.test(storedSession.sessionId())) 483 return false; 484 485 Duration idleDuration = between(storedSession.lastActivityAt(), Instant.now()); 486 return idleDuration.compareTo(this.idleTimeout) > 0; 487 } 488 489 @NonNull 490 private Duration sweepInterval() { 491 return this.idleTimeout.compareTo(DEFAULT_SWEEP_INTERVAL) < 0 ? this.idleTimeout : DEFAULT_SWEEP_INTERVAL; 492 } 493}