Skip to content

Commit 5e36d5f

Browse files
naji247naji247
authored
fix: improve session management (#9)
1 parent 4b64244 commit 5e36d5f

File tree

9 files changed

+1117
-149
lines changed

9 files changed

+1117
-149
lines changed

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "mcpcat",
3-
"version": "0.1.7",
3+
"version": "0.1.8",
44
"description": "Analytics tool for MCP (Model Context Protocol) servers - tracks tool usage patterns and provides insights",
55
"type": "module",
66
"main": "dist/index.js",

‎src/index.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@ import { writeToLog } from "./modules/logging.js";
1616
import { setupMCPCatTools } from "./modules/tools.js";
1717
import { setupToolCallTracing } from "./modules/tracing.js";
1818
import { getSessionInfo, newSessionId } from "./modules/session.js";
19-
import { setServerTrackingData } from "./modules/internal.js";
19+
import {
20+
setServerTrackingData,
21+
getServerTrackingData,
22+
} from "./modules/internal.js";
2023
import { setupTracking } from "./modules/tracingV2.js";
2124
import { TelemetryManager } from "./modules/telemetry.js";
2225
import { setTelemetryManager } from "./modules/eventQueue.js";
@@ -136,6 +139,15 @@ function track(
136139
: validatedServer
137140
) as MCPServerLike;
138141

142+
// Check if server is already being tracked
143+
const existingData = getServerTrackingData(lowLevelServer);
144+
if (existingData) {
145+
writeToLog(
146+
"[SESSION DEBUG] track() - Server already being tracked, skipping initialization",
147+
);
148+
return validatedServer;
149+
}
150+
139151
// Initialize telemetry if exporters are configured
140152
if (options.exporters) {
141153
const telemetryManager = new TelemetryManager(options.exporters);
@@ -167,6 +179,7 @@ function track(
167179
identify: options.identify,
168180
redactSensitiveInformation: options.redactSensitiveInformation,
169181
},
182+
sessionSource: "mcpcat", // Initially MCPCat-generated, will change to "mcp" if MCP sessionId is provided
170183
};
171184

172185
setServerTrackingData(lowLevelServer, mcpcatData);

‎src/modules/internal.ts‎

Lines changed: 271 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,130 @@
1-
import { MCPCatData, MCPServerLike, UserIdentity } from "../types.js";
1+
import {
2+
MCPCatData,
3+
MCPServerLike,
4+
UserIdentity,
5+
CompatibleRequestHandlerExtra,
6+
UnredactedEvent,
7+
} from "../types.js";
8+
import { PublishEventRequestEventTypeEnum } from "mcpcat-api";
9+
import { publishEvent } from "./eventQueue.js";
10+
import { getMCPCompatibleErrorMessage } from "./compatibility.js";
11+
import { writeToLog } from "./logging.js";
12+
import { INACTIVITY_TIMEOUT_IN_MINUTES } from "./constants.js";
13+
14+
/**
15+
* Simple LRU cache for session identities.
16+
* Prevents memory leaks by capping at maxSize entries.
17+
* This cache persists across server instance restarts.
18+
*/
19+
class IdentityCache {
20+
private cache: Map<string, { identity: UserIdentity; timestamp: number }>;
21+
private maxSize: number;
22+
23+
constructor(maxSize: number = 1000) {
24+
this.cache = new Map();
25+
this.maxSize = maxSize;
26+
}
27+
28+
get(sessionId: string): UserIdentity | undefined {
29+
const entry = this.cache.get(sessionId);
30+
if (entry) {
31+
// Update timestamp on access (LRU behavior)
32+
entry.timestamp = Date.now();
33+
// Move to end (most recently used)
34+
this.cache.delete(sessionId);
35+
this.cache.set(sessionId, entry);
36+
return entry.identity;
37+
}
38+
return undefined;
39+
}
40+
41+
set(sessionId: string, identity: UserIdentity): void {
42+
// Remove if already exists (to re-add at end)
43+
this.cache.delete(sessionId);
44+
45+
// Evict oldest if at capacity
46+
if (this.cache.size >= this.maxSize) {
47+
const oldestKey = this.cache.keys().next().value;
48+
if (oldestKey !== undefined) {
49+
this.cache.delete(oldestKey);
50+
}
51+
}
52+
53+
this.cache.set(sessionId, { identity, timestamp: Date.now() });
54+
}
55+
56+
has(sessionId: string): boolean {
57+
return this.cache.has(sessionId);
58+
}
59+
60+
size(): number {
61+
return this.cache.size;
62+
}
63+
}
64+
65+
// Global identity cache shared across all server instances
66+
// This prevents duplicate identify events when server objects are recreated
67+
const _globalIdentityCache = new IdentityCache(1000);
68+
69+
/**
70+
* Maps userId to recent session IDs for reconnection support.
71+
* When a user reconnects (new initialize without MCP sessionId),
72+
* we can reuse their previous session if it's recent enough.
73+
*/
74+
class UserSessionCache {
75+
private cache: Map<string, { sessionId: string; lastSeen: number }>;
76+
private maxSize: number;
77+
78+
constructor(maxSize: number = 1000) {
79+
this.cache = new Map();
80+
this.maxSize = maxSize;
81+
}
82+
83+
getRecentSession(userId: string, timeoutMs: number): string | undefined {
84+
const entry = this.cache.get(userId);
85+
if (!entry) return undefined;
86+
87+
// Check if session has expired
88+
if (Date.now() - entry.lastSeen > timeoutMs) {
89+
this.cache.delete(userId);
90+
return undefined;
91+
}
92+
93+
return entry.sessionId;
94+
}
95+
96+
set(userId: string, sessionId: string): void {
97+
// Remove if already exists (to re-add at end for LRU)
98+
this.cache.delete(userId);
99+
100+
// Evict oldest if at capacity
101+
if (this.cache.size >= this.maxSize) {
102+
const oldestKey = this.cache.keys().next().value;
103+
if (oldestKey !== undefined) {
104+
this.cache.delete(oldestKey);
105+
}
106+
}
107+
108+
this.cache.set(userId, { sessionId, lastSeen: Date.now() });
109+
}
110+
}
111+
112+
// Global user session cache for reconnection support
113+
const _globalUserSessionCache = new UserSessionCache(1000);
114+
115+
/**
116+
* FOR TESTING ONLY: Manually set a user session cache entry with custom lastSeen timestamp
117+
*/
118+
export function _testSetUserSession(
119+
userId: string,
120+
sessionId: string,
121+
lastSeenMs: number,
122+
): void {
123+
(_globalUserSessionCache as any).cache.set(userId, {
124+
sessionId,
125+
lastSeen: lastSeenMs,
126+
});
127+
}
2128

3129
// Internal tracking storage
4130
const _serverTracking = new WeakMap<MCPServerLike, MCPCatData>();
@@ -61,3 +187,147 @@ export function mergeIdentities(
61187
},
62188
};
63189
}
190+
191+
/**
192+
* Handles user identification for a request.
193+
* Calls the identify function if configured, compares with previous identity,
194+
* and publishes an identify event only if the identity has changed.
195+
*
196+
* @param server - The MCP server instance
197+
* @param data - The server tracking data
198+
* @param request - The request object to pass to identify function
199+
* @param extra - Optional extra parameters containing headers, sessionId, etc.
200+
*/
201+
export async function handleIdentify(
202+
server: MCPServerLike,
203+
data: MCPCatData,
204+
request: any,
205+
extra?: CompatibleRequestHandlerExtra,
206+
): Promise<void> {
207+
if (!data.options.identify) {
208+
return;
209+
}
210+
211+
const sessionId = data.sessionId;
212+
let identifyEvent: UnredactedEvent = {
213+
sessionId: sessionId,
214+
resourceName: request.params?.name || "Unknown",
215+
eventType: PublishEventRequestEventTypeEnum.mcpcatIdentify,
216+
parameters: {
217+
request: request,
218+
extra: extra,
219+
},
220+
timestamp: new Date(),
221+
redactionFn: data.options.redactSensitiveInformation,
222+
};
223+
224+
try {
225+
const identityResult = await data.options.identify(request, extra);
226+
if (identityResult) {
227+
// Check for session reconnection (if no MCP sessionId provided in extra)
228+
// If this user had a recent session, switch to it instead of creating new one
229+
if (!extra?.sessionId && identityResult.userId) {
230+
const timeoutMs = INACTIVITY_TIMEOUT_IN_MINUTES * 60 * 1000;
231+
const previousSessionId = _globalUserSessionCache.getRecentSession(
232+
identityResult.userId,
233+
timeoutMs,
234+
);
235+
236+
if (previousSessionId && previousSessionId !== data.sessionId) {
237+
// User has a previous session - reconnect to it
238+
const currentSessionIdentity = _globalIdentityCache.get(
239+
data.sessionId,
240+
);
241+
242+
if (!currentSessionIdentity) {
243+
// Current session is brand new (no identity) - reconnect to previous session
244+
data.sessionId = previousSessionId;
245+
data.lastActivity = new Date();
246+
setServerTrackingData(server, data);
247+
248+
writeToLog(
249+
`Reconnected user ${identityResult.userId} to previous session ${previousSessionId} (current session was new)`,
250+
);
251+
} else if (currentSessionIdentity.userId !== identityResult.userId) {
252+
// Current session belongs to different user - reconnect to user's previous session
253+
data.sessionId = previousSessionId;
254+
data.lastActivity = new Date();
255+
setServerTrackingData(server, data);
256+
257+
writeToLog(
258+
`Reconnected user ${identityResult.userId} to previous session ${previousSessionId}`,
259+
);
260+
}
261+
// If current session already belongs to this user, no need to do anything
262+
} else if (!previousSessionId) {
263+
// User has NO previous session - check if current session belongs to someone else
264+
const currentSessionIdentity = _globalIdentityCache.get(
265+
data.sessionId,
266+
);
267+
if (
268+
currentSessionIdentity &&
269+
currentSessionIdentity.userId !== identityResult.userId
270+
) {
271+
// Current session belongs to different user - create new session
272+
const { newSessionId } = await import("./session.js");
273+
data.sessionId = newSessionId();
274+
data.sessionSource = "mcpcat";
275+
data.lastActivity = new Date();
276+
setServerTrackingData(server, data);
277+
278+
writeToLog(
279+
`Created new session ${data.sessionId} for user ${identityResult.userId} (previous session belonged to ${currentSessionIdentity.userId})`,
280+
);
281+
}
282+
}
283+
}
284+
285+
// Now use the (possibly updated) sessionId for all subsequent operations
286+
const currentSessionId = data.sessionId;
287+
288+
// Check global cache first (works across server instance restarts)
289+
const previousIdentity = _globalIdentityCache.get(currentSessionId);
290+
291+
// Merge identities (overwrite userId/userName, merge userData)
292+
const mergedIdentity = mergeIdentities(previousIdentity, identityResult);
293+
294+
// Only publish if identity has changed
295+
const hasChanged =
296+
!previousIdentity ||
297+
!areIdentitiesEqual(previousIdentity, mergedIdentity);
298+
299+
// Update BOTH caches to keep them in sync
300+
// Global cache: persists across server instances
301+
_globalIdentityCache.set(currentSessionId, mergedIdentity);
302+
// Per-server cache: used by getSessionInfo() for fast local access
303+
data.identifiedSessions.set(data.sessionId, mergedIdentity);
304+
305+
// Track userId → sessionId mapping for reconnection support
306+
_globalUserSessionCache.set(mergedIdentity.userId, currentSessionId);
307+
308+
if (hasChanged) {
309+
writeToLog(
310+
`Identified session ${currentSessionId} with identity: ${JSON.stringify(mergedIdentity)}`,
311+
);
312+
publishEvent(server, identifyEvent);
313+
}
314+
} else {
315+
writeToLog(
316+
`Warning: Supplied identify function returned null for session ${sessionId}`,
317+
);
318+
}
319+
} catch (error) {
320+
writeToLog(
321+
`Warning: Supplied identify function threw an error while identifying session ${sessionId} - ${error}`,
322+
);
323+
identifyEvent.duration =
324+
(identifyEvent.timestamp &&
325+
new Date().getTime() - identifyEvent.timestamp.getTime()) ||
326+
undefined;
327+
identifyEvent.isError = true;
328+
identifyEvent.error = {
329+
message: getMCPCompatibleErrorMessage(error),
330+
};
331+
publishEvent(server, identifyEvent);
332+
}
333+
}

0 commit comments

Comments
 (0)