Skip to content

Commit 8aaa7f7

Browse files
dgellowdgellow
committed
Add trace logging for session lifecycle
1 parent de7cb68 commit 8aaa7f7

File tree

3 files changed

+104
-20
lines changed

3 files changed

+104
-20
lines changed

‎internal/client/session_manager.go

Lines changed: 87 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -136,30 +136,37 @@ func (sm *StdioSessionManager) GetSession(key SessionKey) (*StdioSession, bool)
136136

137137
if ok {
138138
now := time.Now()
139+
lastAccessed := session.lastAccessed.Load()
139140
session.lastAccessed.Store(&now)
140141

142+
internal.LogTraceWithFields("session_manager", "Session accessed", map[string]interface{}{
143+
"sessionID": key.SessionID,
144+
"server": key.ServerName,
145+
"user": key.UserEmail,
146+
"lastAccessed": lastAccessed,
147+
"timeSinceAccess": now.Sub(*lastAccessed).String(),
148+
})
149+
141150
select {
142151
case <-session.ctx.Done():
143152
// Process died, remove it
153+
internal.LogTraceWithFields("session_manager", "Session context cancelled, removing", map[string]interface{}{
154+
"sessionID": key.SessionID,
155+
"server": key.ServerName,
156+
"user": key.UserEmail,
157+
})
144158
sm.RemoveSession(key)
145159
return nil, false
146160
default:
147161
return session, true
148162
}
149163
}
150164

151-
// Debug: log all sessions when not found
152-
sm.mu.RLock()
153-
internal.LogWarnWithFields("session_manager", "Session not found", map[string]interface{}{
154-
"looking_for": key,
155-
"total_sessions": len(sm.sessions),
165+
internal.LogTraceWithFields("session_manager", "Session not found", map[string]interface{}{
166+
"sessionID": key.SessionID,
167+
"server": key.ServerName,
168+
"user": key.UserEmail,
156169
})
157-
for k := range sm.sessions {
158-
internal.LogWarnWithFields("session_manager", "Existing session", map[string]interface{}{
159-
"key": k,
160-
})
161-
}
162-
sm.mu.RUnlock()
163170

164171
return nil, false
165172
}
@@ -170,7 +177,13 @@ func (sm *StdioSessionManager) RemoveSession(key SessionKey) {
170177
session, ok := sm.sessions[key]
171178
if ok {
172179
delete(sm.sessions, key)
180+
internal.LogTraceWithFields("session_manager", "Removing session from map", map[string]interface{}{
181+
"sessionID": key.SessionID,
182+
"server": key.ServerName,
183+
"user": key.UserEmail,
184+
})
173185
}
186+
sessionCount := len(sm.sessions)
174187
sm.mu.Unlock()
175188

176189
if ok {
@@ -192,6 +205,16 @@ func (sm *StdioSessionManager) RemoveSession(key SessionKey) {
192205
"server": key.ServerName,
193206
"user": key.UserEmail,
194207
})
208+
209+
internal.LogTraceWithFields("session_manager", "Session removed with details", map[string]interface{}{
210+
"sessionID": key.SessionID,
211+
"server": key.ServerName,
212+
"user": key.UserEmail,
213+
"created": session.created,
214+
"duration": time.Since(session.created).String(),
215+
"lastAccessed": session.lastAccessed.Load(),
216+
"remainingSessions": sessionCount,
217+
})
195218
}
196219
}
197220

@@ -261,6 +284,14 @@ func (s *StdioSession) DiscoverAndRegisterCapabilities(
261284
"server": serverName,
262285
})
263286

287+
internal.LogTraceWithFields("client", "Starting capability discovery", map[string]interface{}{
288+
"server": serverName,
289+
"sessionID": session.SessionID(),
290+
"userEmail": userEmail,
291+
"requiresUserToken": requiresToken,
292+
"hasTokenSetup": tokenSetup != nil,
293+
})
294+
264295
// Discover and register tools
265296
if err := s.client.addToolsToServer(ctx, mcpServer, userEmail, requiresToken, tokenStore, serverName, setupBaseURL, tokenSetup, session); err != nil {
266297
return err
@@ -280,6 +311,16 @@ func (s *StdioSession) DiscoverAndRegisterCapabilities(
280311
"userTokenRequired": requiresToken,
281312
})
282313

314+
internal.LogTraceWithFields("client", "Capability discovery completed", map[string]interface{}{
315+
"server": serverName,
316+
"sessionID": session.SessionID(),
317+
"userEmail": userEmail,
318+
"requiresUserToken": requiresToken,
319+
"toolsRegistered": true,
320+
"promptsRegistered": true,
321+
"resourcesRegistered": true,
322+
})
323+
283324
// Start ping task if needed
284325
if s.client.needPing {
285326
go s.client.startPingTask(ctx)
@@ -295,6 +336,11 @@ func (sm *StdioSessionManager) checkUserLimits(userEmail string) error {
295336
return nil
296337
}
297338

339+
if sm.maxPerUser == 0 {
340+
// 0 means unlimited
341+
return nil
342+
}
343+
298344
count := sm.getUserSessionCount(userEmail)
299345
if count >= sm.maxPerUser {
300346
internal.LogWarnWithFields("session_manager", "User session limit exceeded", map[string]interface{}{
@@ -306,6 +352,13 @@ func (sm *StdioSessionManager) checkUserLimits(userEmail string) error {
306352
ErrUserLimitExceeded, userEmail, count, sm.maxPerUser)
307353
}
308354

355+
internal.LogTraceWithFields("session_manager", "User session limit check passed", map[string]interface{}{
356+
"user": userEmail,
357+
"currentSessions": count,
358+
"maxPerUser": sm.maxPerUser,
359+
"remaining": sm.maxPerUser - count,
360+
})
361+
309362
return nil
310363
}
311364

@@ -362,6 +415,16 @@ func (sm *StdioSessionManager) createSession(
362415
"user": key.UserEmail,
363416
})
364417

418+
internal.LogTraceWithFields("session_manager", "Session created with details", map[string]interface{}{
419+
"sessionID": key.SessionID,
420+
"server": key.ServerName,
421+
"user": key.UserEmail,
422+
"timeout": sm.defaultTimeout.String(),
423+
"maxPerUser": sm.maxPerUser,
424+
"cleanupInterval": sm.cleanupInterval.String(),
425+
"totalSessions": len(sm.sessions),
426+
})
427+
365428
return session, nil
366429
}
367430

@@ -389,14 +452,27 @@ func (sm *StdioSessionManager) cleanupTimedOutSessions() {
389452
// Find timed out sessions
390453
sm.mu.RLock()
391454
timedOut := make([]SessionKey, 0)
455+
totalSessions := len(sm.sessions)
456+
activeSessions := 0
392457
for key, session := range sm.sessions {
393458
lastAccessed := session.lastAccessed.Load()
394459
if lastAccessed != nil && now.Sub(*lastAccessed) > sm.defaultTimeout {
395460
timedOut = append(timedOut, key)
461+
} else {
462+
activeSessions++
396463
}
397464
}
398465
sm.mu.RUnlock()
399466

467+
if totalSessions > 0 || len(timedOut) > 0 {
468+
internal.LogTraceWithFields("session_manager", "Session cleanup cycle", map[string]interface{}{
469+
"totalSessions": totalSessions,
470+
"activeSessions": activeSessions,
471+
"timedOutSessions": len(timedOut),
472+
"timeout": sm.defaultTimeout.String(),
473+
})
474+
}
475+
400476
for _, key := range timedOut {
401477
internal.LogInfoWithFields("session_manager", "Removing timed out session", map[string]interface{}{
402478
"sessionID": key.SessionID,

‎internal/client/session_test.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ func TestStdioSession_DiscoverAndRegisterCapabilities(t *testing.T) {
4848
mockClient.On("ListTools", ctx, mock.AnythingOfType("mcp.ListToolsRequest")).
4949
Return(listToolsResult, nil).Once()
5050

51-
mockSession.On("SessionID").Return("test-session-123").Twice()
51+
mockSession.On("SessionID").Return("test-session-123").Times(4)
5252

5353
mockSession.On("SetSessionTools", mock.MatchedBy(func(tools map[string]server.ServerTool) bool {
5454
return len(tools) == 2 &&
@@ -102,6 +102,8 @@ func TestStdioSession_DiscoverAndRegisterCapabilities(t *testing.T) {
102102
ServerInfo: mcp.Implementation{Name: "test-server", Version: "1.0"},
103103
}, nil)
104104

105+
mockSession.On("SessionID").Return("test-session-123")
106+
105107
err := stdioSession.DiscoverAndRegisterCapabilities(
106108
ctx,
107109
mockServer,
@@ -156,7 +158,7 @@ func TestStdioSession_DiscoverAndRegisterCapabilities(t *testing.T) {
156158
mockClient.On("ListTools", ctx, mock.AnythingOfType("mcp.ListToolsRequest")).
157159
Return(listToolsResult, nil).Once()
158160

159-
mockSession.On("SessionID").Return("test-session-456").Twice()
161+
mockSession.On("SessionID").Return("test-session-456").Times(4)
160162

161163
mockSession.On("SetSessionTools", mock.MatchedBy(func(tools map[string]server.ServerTool) bool {
162164
return len(tools) == 1 && tools["tool1"].Tool.Name == "tool1"
@@ -244,7 +246,7 @@ func TestStdioSession_DiscoverAndRegisterCapabilities(t *testing.T) {
244246
mockClient.On("ListTools", ctx, mock.AnythingOfType("mcp.ListToolsRequest")).
245247
Return(nil, errors.New("failed to list tools"))
246248

247-
mockSession.On("SessionID").Return("test-session").Once()
249+
mockSession.On("SessionID").Return("test-session")
248250

249251
err := stdioSession.DiscoverAndRegisterCapabilities(
250252
ctx,

‎internal/server/handler.go

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
4040
// Create session manager for stdio servers with configurable timeouts
4141
sessionTimeout := 5 * time.Minute
4242
cleanupInterval := 1 * time.Minute
43+
maxPerUser := 10
4344

4445
// Use config values if available
4546
if cfg.Proxy.Sessions != nil {
@@ -55,11 +56,12 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
5556
"interval": cleanupInterval,
5657
})
5758
}
59+
maxPerUser = cfg.Proxy.Sessions.MaxPerUser
5860
}
5961

6062
sessionManager := client.NewStdioSessionManager(
6163
client.WithTimeout(sessionTimeout),
62-
client.WithMaxPerUser(10),
64+
client.WithMaxPerUser(maxPerUser),
6365
client.WithCleanupInterval(cleanupInterval),
6466
)
6567

@@ -444,6 +446,15 @@ func handleSessionRegistration(
444446
"user": handler.userEmail,
445447
})
446448

449+
internal.LogTraceWithFields("server", "Session registration started", map[string]interface{}{
450+
"sessionID": session.SessionID(),
451+
"server": handler.h.serverName,
452+
"user": handler.userEmail,
453+
"requiresUserToken": handler.config.RequiresUserToken,
454+
"transportType": handler.config.TransportType,
455+
"command": handler.config.Command,
456+
})
457+
447458
stdioSession, err := sessionManager.GetOrCreateSession(
448459
sessionCtx,
449460
key,
@@ -483,11 +494,6 @@ func handleSessionRegistration(
483494
return
484495
}
485496

486-
// Note: We don't need to store anything in the session anymore
487-
// The stdio client is connected directly to the shared MCP server
488-
// Capabilities are automatically registered on the shared MCP server
489-
490-
// Track session in storage
491497
if handler.userEmail != "" {
492498
if store, ok := handler.h.tokenStore.(storage.Storage); ok {
493499
activeSession := storage.ActiveSession{

0 commit comments

Comments
 (0)