diff --git a/internal/server/mcp.go b/internal/server/mcp.go index d54fec256..b43048f14 100644 --- a/internal/server/mcp.go +++ b/internal/server/mcp.go @@ -494,8 +494,11 @@ func NewMCPProxyServer( hooks.AddOnRegisterSession(func(ctx context.Context, sess mcpserver.ClientSession) { sessionID := sess.SessionID() - // Just log the registration - client info and capabilities will be set by OnAfterInitialize - // This hook is primarily for persistent connections (SSE) to track when the session is registered + // Client info and capabilities are set by OnAfterInitialize. A transport + // can re-register an id that was unregistered without a new initialize + // (a GET stream ending and being re-opened), so revive the soft-closed + // entry to keep its attribution (#1205). + sessionStore.Reopen(sessionID) logger.Info("MCP session registered", zap.String("session_id", sessionID), ) @@ -569,9 +572,12 @@ func NewMCPProxyServer( installCallerInstructionsHooks(hooks, proxyRef.Load) // Add hook to clean up session on disconnect - // NOTE: This hook may NOT be called for Streamable HTTP transport because HTTP is stateless - // and has no persistent connection. For HTTP transport, we rely on inactivity timeout - // cleanup (see runtime.backgroundSessionCleanup). + // NOTE: for Streamable HTTP this hook fires whenever a GET stream ends, not + // only on a real disconnect, and the client may keep using the same session id + // afterwards. So it only soft-closes (SessionStore.RemoveSession): attribution + // is kept and revived by the next activity or register, and closed entries are + // evicted after a TTL. Persisted records are also closed by inactivity timeout + // (see runtime.backgroundSessionCleanup). hooks.AddOnUnregisterSession(func(ctx context.Context, sess mcpserver.ClientSession) { sessionID := sess.SessionID() @@ -579,7 +585,7 @@ func NewMCPProxyServer( zap.String("session_id", sessionID), ) - // Remove session information (closes in storage) + // Soft-close the session (closes in storage, keeps attribution) sessionStore.RemoveSession(sessionID) logger.Info("MCP session unregistered", diff --git a/internal/server/session_store.go b/internal/server/session_store.go index 18d9d8ab2..5d4064f22 100644 --- a/internal/server/session_store.go +++ b/internal/server/session_store.go @@ -93,21 +93,39 @@ type SessionInfo struct { // initialize so a binding change can send tools/list_changed on the right // instance (Spec 108 FR-026, plan D18). Never persisted or serialised. server *mcpserver.MCPServer + + // closed marks a soft-closed session: the transport unregistered it, but the + // entry (client name, workspace, persisted flag, work session) is kept so a + // client that re-registers on the same id without a new initialize keeps its + // attribution (#1205). Cleared by the next activity or register; evicted + // after closedSessionTTL. + closed bool + closedAt time.Time } +// closedSessionTTL is how long a soft-closed session is retained. It matches the +// storage inactivity timeout (runtime.backgroundSessionCleanup). +const closedSessionTTL = 30 * time.Minute + // SessionStore manages MCP session information type SessionStore struct { sessions map[string]*SessionInfo // activeProfiles holds the per-session active profile slug selected via the // set_profile MCP tool (Profiles v2 T2). Kept orthogonal to sessions so a - // re-initialize (SetSession) never clobbers a live selection. Cleared on - // session close (RemoveSession) — covering both the OnUnregisterSession hook - // and the background inactivity cleanup, which both call RemoveSession. + // re-initialize (SetSession) never clobbers a live selection. Kept across + // a soft close (RemoveSession) and dropped with the closed entry on eviction. activeProfiles map[string]string mu sync.RWMutex logger *zap.Logger storageManager *storage.Manager + // now is the clock (test seam for closed-entry eviction); nil means time.Now. + now func() time.Time + // lifecycleMu serializes the storage close/reopen reconciliation so a + // concurrent unregister and re-activity cannot leave the stored status + // disagreeing with the in-memory one. + lifecycleMu sync.Mutex + // profileWriter, when non-nil, replaces the storage write-through of a // changed profile resolution. Test seam only (counts writes). profileWriter func(sessionID, profile, source string) @@ -125,6 +143,81 @@ func NewSessionStore(logger *zap.Logger) *SessionStore { } } +func (s *SessionStore) clock() time.Time { + if s.now != nil { + return s.now() + } + return time.Now() +} + +// evictClosedLocked drops soft-closed sessions older than closedSessionTTL, +// along with their set_profile selection. Caller must hold the write lock. +func (s *SessionStore) evictClosedLocked() { + cutoff := s.clock().Add(-closedSessionTTL) + for id, info := range s.sessions { + if info.closed && info.closedAt.Before(cutoff) { + delete(s.sessions, id) + delete(s.activeProfiles, id) + } + } +} + +// reconcileStorage brings the stored status of a persisted session in line with +// its current in-memory open/closed state. Called only after a transition. +func (s *SessionStore) reconcileStorage(sessionID string) { + s.lifecycleMu.Lock() + defer s.lifecycleMu.Unlock() + + s.mu.RLock() + info, ok := s.sessions[sessionID] + persisted, closed := ok && info.persisted, ok && info.closed + done := (chan struct{})(nil) + if ok { + done = info.persistDone + } + mgr := s.storageManager + s.mu.RUnlock() + + if !persisted || mgr == nil { + return + } + <-done // the row may still be in flight + if closed { + if err := mgr.CloseSession(sessionID); err != nil { + s.logger.Warn("failed to close session in storage", + zap.String("session_id", sessionID), zap.Error(err)) + } + return + } + if err := mgr.ReopenSession(sessionID); err != nil { + s.logger.Debug("failed to reopen session in storage", + zap.String("session_id", sessionID), zap.Error(err)) + } +} + +// reopenLocked clears the soft-close flag. It reports whether the session was +// closed. Caller must hold the write lock. +func (s *SessionStore) reopenLocked(info *SessionInfo) bool { + if !info.closed { + return false + } + info.closed = false + info.closedAt = time.Time{} + return true +} + +// Reopen revives a soft-closed session (the transport re-registered the same id +// without a new initialize). A no-op for unknown or already-open sessions. +func (s *SessionStore) Reopen(sessionID string) { + s.mu.Lock() + info, ok := s.sessions[sessionID] + revived := ok && s.reopenLocked(info) + s.mu.Unlock() + if revived { + s.reconcileStorage(sessionID) + } +} + // SetStorageManager sets the storage manager for persistence func (s *SessionStore) SetStorageManager(manager *storage.Manager) { s.mu.Lock() @@ -162,6 +255,7 @@ func (s *SessionStore) SetSession(sessionID, clientName, clientVersion string, h close(info.workspaceReady) } s.sessions[sessionID] = info + s.evictClosedLocked() s.logger.Debug("session registered (not yet persisted — awaiting first activity)", zap.String("session_id", sessionID), @@ -226,10 +320,14 @@ func (s *SessionStore) EnsurePersisted(sessionID string, resolveWorkSession func return "" } + revived := s.reopenLocked(info) if info.persisted { workSessionID := info.workSessionID done := info.persistDone s.mu.Unlock() + if revived { + s.reconcileStorage(sessionID) + } <-done // the row may still be in flight; stats callers need it to exist return workSessionID } @@ -385,31 +483,38 @@ func (s *SessionStore) AbandonWorkspaceFetch(sessionID string) { } } -// RemoveSession removes session information. +// RemoveSession soft-closes a session (#1205). +// +// The entry is KEPT, marked closed: a client that re-registers on the same id +// (e.g. a GET stream ending and being re-opened) sends no new initialize, so +// dropping the entry would lose its client name, workspace, persisted flag and +// work session for good. The next activity or Reopen revives it; otherwise it is +// evicted after closedSessionTTL. // // Only sessions that were actually persisted are closed in storage. A session // that never did any work has no stored record (Spec 082), and asking storage to -// close it would return "session not found" on every idle disconnect — trading -// a flood of junk records for a flood of junk warnings. +// close it would return "session not found" on every idle disconnect. func (s *SessionStore) RemoveSession(sessionID string) { s.mu.Lock() info, ok := s.sessions[sessionID] wasPersisted := ok && info.persisted - delete(s.sessions, sessionID) - delete(s.activeProfiles, sessionID) - mgr := s.storageManager + transition := ok && !info.closed + if ok { + info.closed = true + if transition { + info.closedAt = s.clock() + } + } else { + delete(s.activeProfiles, sessionID) + } + s.evictClosedLocked() s.mu.Unlock() - if wasPersisted && mgr != nil { - if err := mgr.CloseSession(sessionID); err != nil { - s.logger.Warn("failed to close session in storage", - zap.String("session_id", sessionID), - zap.Error(err), - ) - } + if transition && wasPersisted { + s.reconcileStorage(sessionID) } - s.logger.Debug("session info removed", + s.logger.Debug("session soft-closed", zap.String("session_id", sessionID), zap.Bool("was_persisted", wasPersisted), ) @@ -422,13 +527,17 @@ func (s *SessionStore) RemoveSession(sessionID string) { // until the session does its first piece of work. We skip the write for a // never-persisted session rather than logging a warning for it. func (s *SessionStore) UpdateSessionStats(sessionID string, tokens int) { - s.mu.RLock() - persisted := false + s.mu.Lock() + persisted, revived := false, false if info, ok := s.sessions[sessionID]; ok { persisted = info.persisted + revived = s.reopenLocked(info) } mgr := s.storageManager - s.mu.RUnlock() + s.mu.Unlock() + if revived { + s.reconcileStorage(sessionID) + } if !persisted || mgr == nil { return @@ -447,13 +556,17 @@ func (s *SessionStore) UpdateSessionStats(sessionID string, tokens int) { // not gain one here, since merely exchanging MCP messages is not "work" // (Spec 082): that is exactly what the handshake-only agents do. func (s *SessionStore) UpdateActivity(sessionID string) { - s.mu.RLock() - persisted := false + s.mu.Lock() + persisted, revived := false, false if info, ok := s.sessions[sessionID]; ok { persisted = info.persisted + revived = s.reopenLocked(info) } mgr := s.storageManager - s.mu.RUnlock() + s.mu.Unlock() + if revived { + s.reconcileStorage(sessionID) + } if !persisted || mgr == nil { return @@ -527,7 +640,7 @@ func (s *SessionStore) SessionsMatching(pred func(*SessionInfo) bool) []sessionT defer s.mu.RUnlock() var out []sessionTarget for id, info := range s.sessions { - if pred(info) { + if !info.closed && pred(info) { out = append(out, sessionTarget{ID: id, Server: info.server}) } } @@ -595,7 +708,7 @@ func (s *SessionStore) NotifyTargets(tokenName string) []sessionTarget { defer s.mu.RUnlock() var out []sessionTarget for id, info := range s.sessions { - if info.TokenName == tokenName { + if !info.closed && info.TokenName == tokenName { out = append(out, sessionTarget{ID: id, Server: info.server}) } } @@ -676,7 +789,7 @@ func (s *SessionStore) SessionsForToken(tokenName string) []string { defer s.mu.RUnlock() var ids []string for id, info := range s.sessions { - if info.TokenName == tokenName { + if !info.closed && info.TokenName == tokenName { ids = append(ids, id) } } @@ -694,12 +807,18 @@ func (s *SessionStore) GetActiveProfile(sessionID string) string { return s.activeProfiles[sessionID] } -// Count returns the number of active sessions +// Count returns the number of open (not soft-closed) sessions func (s *SessionStore) Count() int { s.mu.RLock() defer s.mu.RUnlock() - return len(s.sessions) + n := 0 + for _, info := range s.sessions { + if !info.closed { + n++ + } + } + return n } // workspaceDisplayName is the basename of a workspace root. Only the basename is diff --git a/internal/server/session_store_softclose_test.go b/internal/server/session_store_softclose_test.go new file mode 100644 index 000000000..d36a7d607 --- /dev/null +++ b/internal/server/session_store_softclose_test.go @@ -0,0 +1,115 @@ +package server + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/zap" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" +) + +// #1205: an unregister followed by a re-register on the same id (no new +// initialize) must keep client attribution and the cached work session. +func TestSessionStore_SoftCloseKeepsAttributionAcrossReRegister(t *testing.T) { + mgr, err := storage.NewManager(t.TempDir(), zap.NewNop().Sugar()) + require.NoError(t, err) + t.Cleanup(func() { _ = mgr.Close() }) + + store := NewSessionStore(zap.NewNop()) + store.SetStorageManager(mgr) + store.SetSession("s1", "claude-code", "1.2", false, false, nil) + + ws := store.EnsurePersisted("s1", func(*SessionInfo) string { return "ws-1" }) + require.Equal(t, "ws-1", ws) + + store.RemoveSession("s1") + rec, err := mgr.GetSessionByID("s1") + require.NoError(t, err) + assert.Equal(t, "closed", rec.Status) + assert.Zero(t, store.Count(), "soft-closed session is not live") + + // Attribution survives the unregister. + require.NotNil(t, store.GetSession("s1")) + assert.Equal(t, "claude-code", store.GetSession("s1").ClientName) + assert.Equal(t, "ws-1", store.WorkSessionID("s1")) + + // Next activity revives it: same work session, no duplicate record, row active. + store.UpdateActivity("s1") + assert.Equal(t, 1, store.Count()) + assert.Equal(t, "ws-1", store.EnsurePersisted("s1", func(*SessionInfo) string { return "other" })) + rec, err = mgr.GetSessionByID("s1") + require.NoError(t, err) + assert.Equal(t, "active", rec.Status) + assert.Nil(t, rec.EndTime) + assert.Equal(t, "claude-code", rec.ClientName) +} + +func TestSessionStore_ReopenOnRegisterHook(t *testing.T) { + store := NewSessionStore(zap.NewNop()) + store.SetSession("s1", "codex", "1", false, false, nil) + store.SetActiveProfile("s1", "work") + store.RemoveSession("s1") + assert.Equal(t, "work", store.GetActiveProfile("s1"), "set_profile selection survives a soft close") + + store.Reopen("s1") + assert.Equal(t, 1, store.Count()) + assert.Equal(t, "codex", store.GetSession("s1").ClientName) + store.Reopen("unknown") // no-op +} + +func TestSessionStore_ClosedEntriesEvictedAfterTTL(t *testing.T) { + store := NewSessionStore(zap.NewNop()) + now := time.Now() + store.now = func() time.Time { return now } + + store.SetSession("s1", "a", "1", false, false, nil) + store.SetActiveProfile("s1", "p") + store.RemoveSession("s1") + + now = now.Add(closedSessionTTL - time.Minute) + store.SetSession("s2", "b", "1", false, false, nil) // triggers sweep + assert.NotNil(t, store.GetSession("s1"), "kept before the TTL") + + now = now.Add(2 * time.Minute) + store.SetSession("s3", "c", "1", false, false, nil) // triggers sweep + assert.Nil(t, store.GetSession("s1"), "evicted after the TTL") + assert.Empty(t, store.GetActiveProfile("s1")) + assert.NotNil(t, store.GetSession("s2")) +} + +func TestSessionStore_RemoveRacesUpdateActivity(t *testing.T) { + store := NewSessionStore(zap.NewNop()) + store.SetSession("s1", "claude-code", "1", false, false, nil) + + var wg sync.WaitGroup + wg.Add(3) + go func() { + defer wg.Done() + for i := 0; i < 300; i++ { + store.RemoveSession("s1") + } + }() + go func() { + defer wg.Done() + for i := 0; i < 300; i++ { + store.UpdateActivity("s1") + store.UpdateSessionStats("s1", 1) + } + }() + go func() { + defer wg.Done() + for i := 0; i < 300; i++ { + store.Reopen("s1") + _ = store.GetSession("s1") + _ = store.Count() + } + }() + wg.Wait() + + require.NotNil(t, store.GetSession("s1")) + assert.Equal(t, "claude-code", store.GetSession("s1").ClientName) +} diff --git a/internal/server/session_store_worksession_test.go b/internal/server/session_store_worksession_test.go index 7ad097e43..40f645e14 100644 --- a/internal/server/session_store_worksession_test.go +++ b/internal/server/session_store_worksession_test.go @@ -28,7 +28,7 @@ func TestSessionStore_HandshakeOnlyIsNotPersisted(t *testing.T) { "a session that never worked has no work session") store.RemoveSession("s1") - assert.Nil(t, store.GetSession("s1")) + assert.Zero(t, store.Count(), "a soft-closed session is not live") } // GetSession must hand back a COPY. The roots goroutine writes Workspace under diff --git a/internal/storage/manager.go b/internal/storage/manager.go index c4ef59f62..bed16db9b 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -1549,6 +1549,10 @@ func (m *Manager) CreateSession(session *SessionRecord) error { existingSession.HasRoots = session.HasRoots existingSession.HasSampling = session.HasSampling existingSession.Experimental = session.Experimental + // A re-created session is live again (re-initialize on a + // soft-closed id must not leave the record closed). + existingSession.Status = "active" + existingSession.EndTime = nil session = &existingSession m.logger.Debugw("Updating existing session with new data", "session_id", session.ID, "client_name", session.ClientName) break @@ -1631,6 +1635,42 @@ func (m *Manager) CloseSession(sessionID string) error { }) } +// ReopenSession marks a closed session active again and clears its end time. It +// is a no-op for an unknown or already-active session (no duplicate record). +func (m *Manager) ReopenSession(sessionID string) error { + m.mu.Lock() + defer m.mu.Unlock() + + return m.db.db.Update(func(tx *bbolt.Tx) error { + bucket := tx.Bucket([]byte(SessionsBucket)) + if bucket == nil { + return nil + } + c := bucket.Cursor() + for k, v := c.First(); k != nil; k, v = c.Next() { + if !strings.HasSuffix(string(k), "_"+sessionID) { + continue + } + var session SessionRecord + if err := json.Unmarshal(v, &session); err != nil { + return fmt.Errorf("failed to unmarshal session: %w", err) + } + if session.Status == "active" { + return nil + } + session.Status = "active" + session.EndTime = nil + session.LastActivity = time.Now() + data, err := json.Marshal(session) + if err != nil { + return fmt.Errorf("failed to marshal session: %w", err) + } + return bucket.Put(k, data) + } + return nil + }) +} + // GetRecentSessions returns the sessions that were active most recently. // // "Recent" means LastActivity, not StartTime. Bucket keys are